//! # Burn
//!
//! Burn is a new comprehensive dynamic Deep Learning Framework built using Rust
//! with extreme flexibility, compute efficiency and portability as its primary goals.
//!
//! ## Performance
//!
//! Because we believe the goal of a deep learning framework is to convert computation
//! into useful intelligence, we have made performance a core pillar of Burn.
//! We strive to achieve top efficiency by leveraging multiple optimization techniques:
//!
//! - Automatic kernel fusion
//! - Asynchronous execution
//! - Thread-safe building blocks
//! - Intelligent memory management
//! - Automatic kernel selection
//! - Hardware specific features
//! - Custom Backend Extension
//!
//! ## Training & Inference
//!
//! The whole deep learning workflow is made easy with Burn, as you can monitor your training progress
//! with an ergonomic dashboard, and run inference everywhere from embedded devices to large GPU clusters.
//!
//! Burn was built from the ground up with training and inference in mind. It's also worth noting how Burn,
//! in comparison to frameworks like PyTorch, simplifies the transition from training to deployment,
//! eliminating the need for code changes.
//!
//! ## Backends
//!
//! Burn strives to be as fast as possible on as many hardwares as possible, with robust implementations.
//! We believe this flexibility is crucial for modern needs where you may train your models in the cloud,
//! then deploy on customer hardwares, which vary from user to user.
//!
//! Compared to other frameworks, Burn has a very different approach to supporting many backends.
//! By design, most code is generic over the Backend trait, which allows us to build Burn with swappable backends.
//! This makes composing backend possible, augmenting them with additional functionalities such as
//! autodifferentiation and automatic kernel fusion.
//!
//! - WGPU (WebGPU): Cross-Platform GPU Backend
//! - Candle: Backend using the Candle bindings
//! - LibTorch: Backend using the LibTorch bindings
//! - NdArray: Backend using the NdArray primitive as data structure
//! - Autodiff: Backend decorator that brings backpropagation to any backend
//! - Fusion: Backend decorator that brings kernel fusion to backends that support it
//!
//! # Quantization (Beta)
//!
//! Quantization techniques perform computations and store tensors in lower precision data types like 8-bit integer
//! instead of floating point precision. There are multiple approaches to quantize a deep learning model. In most cases,
//! the model is trained in floating point precision and later converted to the lower precision data type. This is called
//! post-training quantization (PTQ). On the other hand, quantization aware training (QAT) models the effects of quantization
//! during training. Quantization errors are thus modeled in the forward and backward passes, which helps the model learn
//! representations that are more robust to the reduction in precision.
//!
//! Quantization support in Burn is currently in active development. It supports the following modes on some backends:
//! - Per-tensor and per-block (linear) quantization to 8-bit, 4-bit and 2-bit representations
//!
//! ## Feature Flags
//!
//! The following feature flags are available.
//! By default, the feature `std` is activated.
//!
//! - Training
//!   - `train`: Enables features `dataset` and `autodiff` and provides a training environment
//!   - `tui`: Includes Text UI with progress bar and plots
//!   - `metrics`: Includes system info metrics (CPU/GPU usage, etc.)
//! - Dataset
//!   - `dataset`: Includes a datasets library
//!   - `audio`: Enables audio datasets (SpeechCommandsDataset)
//!   - `sqlite`: Stores datasets in SQLite database
//!   - `sqlite_bundled`: Use bundled version of SQLite
//!   - `vision`: Enables vision datasets (MnistDataset)
//! - Backends
//!   - `wgpu`: Makes available the WGPU backend
//!   - `webgpu`: Makes available the `wgpu` backend with the WebGPU Shading Language (WGSL) compiler
//!   - `vulkan`: Makes available the `wgpu` backend with the alternative SPIR-V compiler
//!   - `cuda`: Makes available the CUDA backend
//!   - `rocm`: Makes available the ROCm backend
//!   - `candle`: Makes available the Candle backend
//!   - `tch`: Makes available the LibTorch backend
//!   - `ndarray`: Makes available the NdArray backend
//! - Backend specifications
//!   - `accelerate`: If supported, Accelerate will be used
//!   - `blas-netlib`: If supported, Blas Netlib will be use
//!   - `openblas`: If supported, Openblas will be use
//!   - `openblas-system`: If supported, Openblas installed on the system will be use
//!   - `autotune`: Enable running benchmarks to select the best kernel in backends that support it.
//!   - `fusion`: Enable operation fusion in backends that support it.
//! - Backend decorators
//!   - `autodiff`: Makes available the Autodiff backend
//! - Model Storage
//!   - `store`: Enables model storage with SafeTensors format and PyTorch interoperability
//! - Others:
//!   - `std`: Activates the standard library (deactivate for no_std)
//!   - `server`: Enables the remote server.
//!   - `network`: Enables network utilities (currently, only a file downloader with progress bar)
//!   - `experimental-named-tensor`: Enables named tensors (experimental)
//!
//! You can also check the details in sub-crates [`burn-core`](https://docs.rs/burn-core) and [`burn-train`](https://docs.rs/burn-train).
/// Train module
/// Backend module.
/// Module for collective operations
/// Module for model storage and serialization
/// Neural network module.
/// Optimizers module.
// For backward compat, `burn::lr_scheduler::*`
/// Learning rate scheduler module.
// For backward compat, `burn::grad_clipping::*`
/// Gradient clipping module.
/// CubeCL module re-export.
//! Structs and macros used by most projects. Add `use
//! burn::prelude::*` to your code to quickly get started with
//! Burn.
/// Enable auto-differentiation on a backend.
///
/// This works as a backend decorator, extending the functionality of any backend with
/// backpropagation.
pub struct Autodiff<B, C = NoCheckpointing> {
impl<B: Backend, C: CheckpointStrategy> Backend for Autodiff<B, C> {
    type Device = B::Device;
    type FloatTensorPrimitive = AutodiffTensor<B>;
    type FloatElem = B::FloatElem;
    type IntTensorPrimitive = B::IntTensorPrimitive;
    type IntElem = B::IntElem;
    type BoolTensorPrimitive = B::BoolTensorPrimitive;
    type BoolElem = B::BoolElem;
    type QuantizedTensorPrimitive = B::QuantizedTensorPrimitive;
    type QuantizedEncoding = B::QuantizedEncoding;
    fn ad_enabled() -> bool {
    fn name(device: &Self::Device) -> String {
    fn seed(seed: u64) {
    fn sync(device: &B::Device) {
impl<B: Backend, C: CheckpointStrategy> AutodiffBackend for Autodiff<B, C> {
    type InnerBackend = B;
    type Gradients = Gradients;
    fn backward(tensor: AutodiffTensor<B>) -> Gradients {
    fn grad(tensor: &AutodiffTensor<B>, grads: &Gradients) -> Option<B::FloatTensorPrimitive> {
    fn grad_remove(
        tensor: &AutodiffTensor<B>,
        grads: &mut Gradients,
    ) -> Option<B::FloatTensorPrimitive> {
    fn inner(tensor: AutodiffTensor<B>) -> B::FloatTensorPrimitive {
    fn from_inner(tensor: B::FloatTensorPrimitive) -> AutodiffTensor<B> {
    fn grad_replace(
        tensor: &AutodiffTensor<B>,
        grads: &mut Self::Gradients,
        grad: B::FloatTensorPrimitive,
    ) {
    fn int_inner(tensor: IntTensor<Self>) -> IntTensor<Self::InnerBackend> {
    fn bool_inner(tensor: BoolTensor<Self>) -> BoolTensor<Self::InnerBackend> {
    fn int_from_inner(tensor: IntTensor<Self::InnerBackend>) -> IntTensor<Self> {
    fn bool_from_inner(tensor: BoolTensor<Self::InnerBackend>) -> BoolTensor<Self> {
    fn q_inner(tensor: QuantizedTensor<Self>) -> QuantizedTensor<Self::InnerBackend> {
    fn q_from_inner(tensor: QuantizedTensor<Self::InnerBackend>) -> QuantizedTensor<Self> {
/// Links a [NodeID] to its autodiff graph [NodeRef]
pub(crate) struct NodeTree {
impl NodeTree {
/// Gives the parents of the node in the autodiff graph
    pub(crate) fn parents(&self, node_id: &NodeID) -> Option<Vec<NodeID>> {
/// Struct responsible of fetching the output for a node in the autodiff graph during a backward pass
pub struct Checkpointer {
impl Checkpointer {
/// Gives the output of the given node, by recursively asking parents to compute themselves
/// or give their pre-computed tensors.
    pub fn retrieve_node_output<T>(&mut self, node_id: NodeID) -> T
    where
        T: Clone + Send + 'static,
    {
/// Sorts the ancestors of NodeID in a way such that all parents come before their children
/// Useful to avoid recursivity later when mutating the states
///
/// The sort on a compute bound state or a memory bound that is already computed is trivial.
/// The match on State::Computed also serves as a stopping criterion for the sort,
/// we don't need to look higher than that during recursivity.
    fn topological_sort(&self, node_id: NodeID) -> Vec<NodeID> {
/// Checks if checkpointer has been drained adequately. Useful for testing
    pub fn is_empty(&self) -> bool {
/// Determines if a node should checkpoint its computed output or its retro_forward for recomputation
/// The action is normally created by the child of the node, once the node is determined to be needed
pub enum CheckpointingAction {
/// The node's already computed output should be saved
/// The node
/// The node's output
/// The node should recompute itself when asked
/// The node
/// How the node should recompute itself
// TODO: Remove that when proper client server.
unsafe impl Send for CheckpointingAction {
impl CheckpointingAction {
/// Utilitary function to access the id of the node of the checkpointing action
    pub fn id(&self) -> NodeID {
/// Accumulates checkpoints as checkpointing actions during the forward pass,
/// and builds a checkpointer right before the backward pass
pub struct CheckpointerBuilder {
/// Determines if a checkpoint should impact the n_required values (Main)
/// or if it should just keep the state in case it's required (Backup)
///
pub(crate) enum ActionType {
/// Explicit actions have been explicitly requested by some operation to retrieve their state
/// Backup actions are not always needed. They exist to save the output of an operation
/// whose child is memory bound, in case the state is indirectly needed when computing
/// the child's retro_forward. If no explicit action ever asks for the child's output, then
/// the backup output will go out of scope when the checkpointer is built.
impl CheckpointerBuilder {
    pub(crate) fn checkpoint<B: Backend>(
        &mut self,
        tensor: &AutodiffTensor<B>,
        action_type: ActionType,
    ) {
    pub(crate) fn extend(&mut self, other: CheckpointerBuilder) {
    pub(crate) fn build(self, node_tree: NodeTree) -> Checkpointer {
// Find recursion stopping points
// We start by identifying how many times each node will be required.
// Then we checkpoint the nodes with the corresponding n_required value
    fn find_stop_nodes(&self) -> Vec<NodeID> {
    fn build_n_required_map(
        &self,
        node_tree: &NodeTree,
        stop_nodes: Vec<NodeID>,
    ) -> HashMap<NodeID, usize> {
    fn insert_checkpoints(
        mut self,
        backward_states_map: &mut HashMap<NodeID, State>,
        retro_forward_map: &mut HashMap<NodeID, Arc<dyn RetroForward>>,
        n_required_map: HashMap<NodeID, usize>,
    ) {
// We do not loop over checkpointing actions anymore because they can contain
// duplicates or miss some that are in backup. We loop over the n_required_map
// from which we use the ids to find them again in the checkpointing actions
// We find the checkpointing action for node_id. It's likely in checkpointing_actions
// so we check there first, otherwise it will be in backup.
// Technically it can be there several times but can never be of both types, so we can assume the first we find is fine
    fn update_n_required_of_parents(
        id: NodeID,
        n_required_map: &mut HashMap<NodeID, usize>,
        node_tree: &NodeTree,
        stop_nodes: &Vec<NodeID>,
    ) {
    fn checkpoint_compute(
        &self,
        backward_states_map: &mut HashMap<NodeID, State>,
        node_id: NodeID,
        state_content: Box<dyn Any + Send>,
        n_required: usize,
    ) {
    fn checkpoint_lazy(
        &self,
        backward_states_map: &mut HashMap<NodeID, State>,
        retro_forward_map: &mut HashMap<NodeID, Arc<dyn RetroForward>>,
        node_id: NodeID,
        retro_forward: Arc<dyn RetroForward>,
        n_required: usize,
    ) {
/// Checkpointer module
/// RetroForward module
/// BackwardStates module
/// CheckpointStrategy module
/// Definition of the forward function of a node, called during retropropagation only.
/// This is different from the normal forward function because it reads and writes from
/// the [BackwardStates] map instead of having a clear function signature.
pub trait RetroForward: Debug + Send + 'static {
/// Applies the forward pass for retropropagation.
    fn forward(&self, states: &mut BackwardStates, out_node: NodeID);
/// Links [NodeID]s to their corresponding [RetroForward]
pub(crate) struct RetroForwards {
impl RetroForwards {
/// Executes the [RetroForward] for a given [NodeID] if the node's
/// [State] is [State::Recompute], otherwise does nothing.
    pub(crate) fn execute_retro_forward(
        &mut self,
        node_id: NodeID,
        backward_states: &mut BackwardStates,
    ) {
// Retro forwards are always used only once because afterwards their state is computed
    pub(crate) fn is_empty(&self) -> bool {
/// Creates a RetroForward struct for unary scalar operations
        struct $name<B: Backend> {
        impl<B: Backend> RetroForward for $name<B> {
            fn forward(&self, states: &mut BackwardStates, out_node: NodeID) {
/// Creates a RetroForward struct for unary scalar operations
        struct $name<B: Backend> {
        impl<B: Backend> RetroForward for $name<B> {
            fn forward(&self, states: &mut BackwardStates, out_node: NodeID) {
/// Creates a RetroForward struct for binary operations
        struct $name<B: Backend> {
        impl<B: Backend> RetroForward for $name<B> {
            fn forward(&self, states: &mut BackwardStates, out_node: NodeID) {
/// In order to accept arbitrary node output in the same hashmap, we need to upcast them to any.
pub(crate) type StateContent = Box<dyn Any + Send>;
/// The state contained at one node. Encapsulates the node output if precomputed,
/// or clearly asks that it needs to be recomputed from the parents.
/// Also keeps track of the number of times the state is required so it can be removed
/// from the map of states on its last use.
pub(crate) enum State {
/// The state was not checkpointed, will need to recompute it from the node's parents
/// The state was checkpointed or computed during retropropagation and can be directly accessed
impl State {
/// Returns a reference to the (not yet) downcasted node output, if checkpointed
    pub(crate) fn to_state_content(&self) -> &StateContent {
/// Returns a (not yet) downcasted node output, if checkpointed
    pub(crate) fn into_state_content(self) -> StateContent {
/// Returns the number of time the state is required
    pub(crate) fn n_required(&self) -> usize {
/// Links [NodeID]s to their current state
pub struct BackwardStates {
impl BackwardStates {
/// Returns the output in the state of the given [NodeID],
/// and decrements the number of times this state is required.
/// This function always gives ownership of the output, but will clone it if needed for further uses.
    pub fn get_state<T>(&mut self, node_id: &NodeID) -> T
    where
        T: Clone + Send + 'static,
    {
// Fetch the state and decrement its number of required
// Downcast the state to whatever it is supposed to be
// If still needed after giving ownership, we copy it back to the hashmap
/// Returns a reference to the [State] of the given node
/// Useful when we need [State] information without needing the underlying tensor
    pub(crate) fn get_state_ref(&self, node_id: &NodeID) -> Option<&State> {
/// Associates a [State] to its [NodeID]
    pub(crate) fn insert_state(&mut self, node_id: NodeID, state: State) {
/// Saves the output to the state of the given [NodeID].
    pub fn save<T>(&mut self, node_id: NodeID, saved_output: T)
    where
        T: Clone + Send + 'static,
    {
    pub(crate) fn is_empty(&self) -> bool {
/// Strategy for the amount of checkpointing to do during autodiff
pub trait CheckpointStrategy: Clone + Copy + Debug + Default + Send + Sync + 'static {
/// May modify the compute property depending on the strategy
    fn compute_property<R: RetroForward>(retro_forward: R) -> ComputingProperty;
/// Checkpoints parents if necessary in the strategy
    fn checkpoint_parents<'a, B2, A>(
        parents: A,
        builder: &mut CheckpointerBuilder,
    ) -> Result<(), CheckpointingError>
    where
        B2: Backend,
        A: IntoIterator<Item = &'a AutodiffTensor<B2>>;
/// Error that can happen when trying to checkpoint a tensor.
pub enum CheckpointingError {
/// When a parent is untracked, we can't easily checkpoint its state, since we don't know the
/// requirements in advanced.
/// All operations are considered compute bound, notwithstanding how they are marked
pub struct NoCheckpointing {
impl CheckpointStrategy for NoCheckpointing {
/// An operation marked as memory bound is actually compute bound.
    fn compute_property<R: RetroForward>(_retro_forward: R) -> ComputingProperty {
/// An operation marked as memory bound is actually compute bound.
/// It's therefore useless to checkpoint the parents
    fn checkpoint_parents<'a, B2, A>(
        _parents: A,
        _builder: &mut CheckpointerBuilder,
    ) -> Result<(), CheckpointingError>
    where
        B2: Backend,
        A: IntoIterator<Item = &'a AutodiffTensor<B2>>,
    {
// Nothing to do here
/// Operation properties are as they are marked (compute or memory bound)
pub struct BalancedCheckpointing {
impl CheckpointStrategy for BalancedCheckpointing {
/// An operation marked as memory bound is memory bound.
/// When memory bound, an operation needs to save its RetroForward
    fn compute_property<R: RetroForward>(retro_forward: R) -> ComputingProperty {
/// An operation marked as memory bound is really memory bound.
/// Since the operation may not checkpoint its parents but may need them indirectly
/// if asked to recompute itself, the method needs to know the parent tensors to maybe checkpoint them
    fn checkpoint_parents<'a, B2, A>(
        parents: A,
        builder: &mut CheckpointerBuilder,
    ) -> Result<(), CheckpointingError>
    where
        B2: Backend,
        A: IntoIterator<Item = &'a AutodiffTensor<B2>>,
    {
/// Gradient identifier.
pub type GradID = u64;
/// Gradients container used during the backward pass.
pub struct Gradients {
impl Gradients {
/// Creates a new gradients container.
    pub fn new<B: Backend>(root_node: NodeRef, root_tensor: FloatTensor<B>) -> Self {
/// Consumes the gradients for a given tensor.
///
/// Each tensor should be consumed exactly 1 time if its gradients are only required during the
/// backward pass, otherwise, it may be consume multiple times.
    pub fn consume<B: Backend>(&mut self, node: &NodeRef) -> FloatTensor<B> {
/// Removes a grad tensor from the container.
    pub fn remove<B: Backend>(&mut self, tensor: &AutodiffTensor<B>) -> Option<FloatTensor<B>> {
/// Gets a grad tensor from the container.
    pub fn get<B: Backend>(&self, tensor: &AutodiffTensor<B>) -> Option<FloatTensor<B>> {
/// Register a grad tensor in the container.
///
/// If the tensor already exists, add both tensors together before saving the result.
    pub fn register<B: Backend>(&mut self, node_id: NodeID, value: FloatTensor<B>) {
/// Backward step for reverse mode autodiff.
pub trait Step: Send + core::fmt::Debug {
/// Executes the step and consumes it.
    fn step(self: Box<Self>, grads: &mut Gradients, checkpointer: &mut Checkpointer);
/// Depth of the operation relative to the first node added to a graph.
    fn depth(&self) -> usize;
/// The node associated to the step.
    fn node(&self) -> NodeID;
/// The parents of the node associated to the step.
    fn parents(&self) -> Vec<NodeID>;
pub type StepBoxed = Box<dyn Step>;
pub enum ComputingProperty {
// Maybe autotune someday
/// This is safe only because we only call RetroForward on the autodiff server.
/// Therefore, the trait will never be used by multiple threads at the same time.
///
/// TODO: Find a way to avoid cloning the compute property, which will remove the need to add the
/// Arc, which will make (dyn RetroForward) safely implement Send.
unsafe impl Send for ComputingProperty {
/// unsafe Sync is required because Send is only implemented for Arc<Sync>, not Arc<Send>.
unsafe impl Sync for ComputingProperty {
/// A node contains graph metadata and should be used wrapped in an Arc for cheap cloning.
pub struct Node {
pub type NodeRef = Arc<Node>;
impl Node {
/// Returns the [node](Node) only if gradients are required.
    pub fn clone_if_require_grad(self: &Arc<Self>) -> Option<NodeRef> {
/// Unique identifier generated for each node.
pub struct NodeID {
/// The integer representation of the id
impl NodeID {
/// Create a unique [node id](NodeID).
    pub fn new() -> Self {
        static COUNTER: AtomicU64 = AtomicU64::new(0);
impl Default for NodeID {
    fn default() -> Self {
/// Requirement for each tensor in the graph.
pub enum Requirement {
/// Operations that require gradients.
/// Operations that require gradients only for backprop.
/// Operations that don't need gradients, therefore not to be included in the graph.
impl Requirement {
/// Returns true if gradients are not required.
    pub fn is_none(&self) -> bool {
/// Returns the right requirement from a list of nodes.
    pub fn from_nodes(nodes: &[NodeRef]) -> Self {
    fn infer(&self, other: &Self) -> Self {
/// Breadth for search algorithm.
pub struct BreadthFirstSearch;
pub trait TraversalItem {
    fn id(&self) -> NodeID;
    fn parents(&self) -> Vec<NodeID>;
impl BreadthFirstSearch {
/// Traverse the graph of backward steps from a root node.
    pub fn traverse<F, I>(
        &self,
        root_id: NodeID,
        root_step: I,
        steps: &mut HashMap<NodeID, I>,
        mut callback: F,
    ) where
        F: FnMut(NodeID, I),
        I: TraversalItem,
    {
impl TraversalItem for StepBoxed {
    fn id(&self) -> NodeID {
    fn parents(&self) -> Vec<NodeID> {
//! # Burn Autodiff
//!
//! This autodiff library is a part of the Burn project. It is a standalone crate
//! that can be used to perform automatic differentiation on tensors. It is
//! designed to be used with the Burn Tensor crate, but it can be used with any
//! tensor library that implements the `Backend` trait.
/// Checkpoint module.
/// Gradients module.
/// Operation module.
// Exported for backend extension
/// A facade around for HashMap and HashSet.
/// This avoids elaborate import wrangling having to happen in every module.
impl<B: Backend, C: CheckpointStrategy> ActivationOps<Autodiff<B, C>> for Autodiff<B, C> {
    fn gelu(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
        struct Gelu;
        impl<B: Backend> Backward<B, 1> for Gelu {
            type State = NodeID;
            fn backward(
                self,
                ops: Ops<Self::State, 1>,
                grads: &mut Gradients,
                checkpointer: &mut Checkpointer,
            ) {
    fn relu(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
        struct Relu;
        impl<B: Backend> Backward<B, 1> for Relu {
            type State = NodeID;
            fn backward(
                self,
                ops: Ops<Self::State, 1>,
                grads: &mut Gradients,
                checkpointer: &mut Checkpointer,
            ) {
    fn sigmoid(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
        struct Sigmoid;
        impl<B: Backend> Backward<B, 1> for Sigmoid {
            type State = NodeID;
            fn backward(
                self,
                ops: Ops<Self::State, 1>,
                grads: &mut Gradients,
                checkpointer: &mut Checkpointer,
            ) {
    fn log_sigmoid(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
        struct LogSigmoid;
        impl<B: Backend> Backward<B, 1> for LogSigmoid {
            type State = NodeID;
            fn backward(
                self,
                ops: Ops<Self::State, 1>,
                grads: &mut Gradients,
                checkpointer: &mut Checkpointer,
            ) {
/// Trait for all operations.
///
/// # Notes
///
/// Concrete types implementing this trait should not have any state.
/// If a state is necessary during the backward pass,
/// they should be declared with the associated type 'State'.
pub trait Backward<B, const N: usize>: Send + core::fmt::Debug
where
    Self: Sized + 'static,
    B: Backend,
{
/// Associated type to compute the backward pass.
    type State: Clone + Send + core::fmt::Debug + 'static;
/// The backward pass.
    fn backward(
        self,
        ops: Ops<Self::State, N>,
        grads: &mut Gradients,
        checkpointer: &mut Checkpointer,
    );
/// Prepare the backward ops.
    fn prepare<C: CheckpointStrategy>(
        self,
        nodes: [NodeRef;
// If not specified we start with ambiguous
/// Execute a binary operation during the backward step.
pub fn binary<B, FLhs, FRhs>(
    parents: [Option<NodeRef>;
/// Execute a unary operation during the backward step.
pub fn unary<B, F>(parents: [Option<NodeRef>;
/// Operation in preparation.
///
/// Each mode has its own set of functions to minimize cloning for unused backward states.
pub struct OpsPrep<Backward, B, S, C, const N: usize, Mode = Init> {
/// Operation is initialized
pub struct Init;
/// Operation has been tagged as memory bound
pub struct MemoryBound;
/// Memory bound operation has received its RetroForward
pub struct MemoryBoundRetroForward;
/// Operation's compute property is fixed
pub struct ComputePropertyDone;
/// Tracked operation tag.
pub struct Tracked;
/// Untracked operation tag.
pub struct UnTracked;
impl<BO, B, S, C, const N: usize> OpsPrep<BO, B, S, C, N, Init>
where
    B: Backend,
    BO: Backward<B, N, State = S>,
{
/// Indicates that the operation is compute bound, meaning its computation
/// is heavy and should not be recomputed
    pub fn compute_bound(self) -> OpsPrep<BO, B, S, C, N, ComputePropertyDone> {
/// Indicates that the operation is memory bound, meaning its computation
/// is light and can be recomputed
    pub fn memory_bound(self) -> OpsPrep<BO, B, S, C, N, MemoryBound> {
impl<BO, B, S, C, const N: usize> OpsPrep<BO, B, S, C, N, MemoryBound>
where
    B: Backend,
    BO: Backward<B, N, State = S>,
    C: CheckpointStrategy,
{
/// Registers the retro forward, if needed
    pub fn retro_forward<R: RetroForward>(
        self,
        retro_forward: R,
    ) -> OpsPrep<BO, B, S, C, N, MemoryBoundRetroForward> {
impl<BO, B, S, C, const N: usize> OpsPrep<BO, B, S, C, N, MemoryBoundRetroForward>
where
    B: Backend,
    BO: Backward<B, N, State = S>,
    C: CheckpointStrategy,
{
/// Checkpoints the parents, if needed
    pub fn parents<'a, B2, A>(mut self, parents: A) -> OpsPrep<BO, B, S, C, N, ComputePropertyDone>
    where
        B2: Backend,
        A: IntoIterator<Item = &'a AutodiffTensor<B2>>,
    {
impl<BO, B, C, const N: usize> OpsPrep<BO, B, (), C, N, ComputePropertyDone>
where
    B: Backend,
    BO: Backward<B, N, State = ()>,
{
/// Prepare a stateless operation.
    pub fn stateless(self, output: FloatTensor<B>) -> AutodiffTensor<B> {
impl<BO, B, S, C, const N: usize> OpsPrep<BO, B, S, C, N, ComputePropertyDone>
where
    B: Backend,
    S: Clone + Send + core::fmt::Debug + 'static,
    BO: Backward<B, N, State = S>,
{
/// Prepare an operation that requires a state during the backward pass.
    pub fn stateful(self) -> OpsKind<BO, B, S, C, N> {
impl<BO, B, S, C, const N: usize> OpsPrep<BO, B, S, C, N, UnTracked>
where
    B: Backend,
    S: Clone + Send + core::fmt::Debug + 'static,
    BO: Backward<B, N, State = S>,
{
/// Finish the preparation of an untracked operation and returns the output tensor.
    pub fn finish(self, output: FloatTensor<B>) -> AutodiffTensor<B> {
// We register the ops in the graph even if untracked, otherwise memory bound operations
// that have an untracked parent would not be able to retrieve it
impl<BO, B, S, C, const N: usize> OpsPrep<BO, B, S, C, N, Tracked>
where
    B: Backend,
    S: Clone + Send + core::fmt::Debug + 'static,
    BO: Backward<B, N, State = S>,
{
/// Finish the preparation of a tracked operation and returns the output tensor.
    pub fn finish(self, state: S, output: FloatTensor<B>) -> AutodiffTensor<B> {
/// Checkpoints the tensor
    pub fn checkpoint(&mut self, tensor: &AutodiffTensor<B>) -> NodeID {
/// Enum used before finishing tracked and untracked operations.
pub enum OpsKind<BO, B, S, C, const N: usize> {
/// Tracked operation preparation.
/// Untracked operation preparation.
/// Operation containing its parent nodes, its own node and the backward step state.
pub struct Ops<S, const N: usize> {
/// Parents nodes.
/// The node.
/// The state.
/// Operation implementing backward [step](Step) with type erasing.
struct OpsStep<B, T, SB, const N: usize>
where
    B: Backend,
    T: Backward<B, N, State = SB>,
    SB: Clone + Send + core::fmt::Debug + 'static,
{
impl<B, T, SB, const N: usize> Step for OpsStep<B, T, SB, N>
where
    B: Backend,
    T: Backward<B, N, State = SB>,
    SB: Clone + Send + core::fmt::Debug + 'static,
{
    fn step(self: Box<Self>, grads: &mut Gradients, checkpointer: &mut Checkpointer) {
    fn node(&self) -> NodeID {
    fn parents(&self) -> Vec<NodeID> {
    fn depth(&self) -> usize {
struct UntrackedOpsStep<const N: usize> {
impl<const N: usize> Step for UntrackedOpsStep<N> {
    fn step(self: Box<Self>, _grads: &mut Gradients, _checkpointer: &mut Checkpointer) {
// Nothing to do
    fn node(&self) -> NodeID {
    fn parents(&self) -> Vec<NodeID> {
    fn depth(&self) -> usize {
/// Make sure the grad tensor has the given shape.
///
/// If broadcasting happened during the forward pass, the gradients will be sum along the
/// broadcasted dimension.
pub fn broadcast_shape<B: Backend>(mut grad: FloatTensor<B>, shape: &Shape) -> FloatTensor<B> {
impl<B: Backend, C: CheckpointStrategy> BoolTensorOps<Self> for Autodiff<B, C> {
    fn bool_from_data(data: TensorData, device: &Device<B>) -> BoolTensor<B> {
    async fn bool_into_data(tensor: BoolTensor<B>) -> TensorData {
    fn bool_into_int(tensor: BoolTensor<B>) -> IntTensor<B> {
    fn bool_to_device(tensor: BoolTensor<B>, device: &Device<B>) -> BoolTensor<B> {
    fn bool_device(tensor: &BoolTensor<B>) -> Device<B> {
    fn bool_reshape(tensor: BoolTensor<B>, shape: Shape) -> BoolTensor<B> {
    fn bool_slice(tensor: BoolTensor<B>, ranges: &[core::ops::Range<usize>]) -> BoolTensor<B> {
    fn bool_empty(shape: Shape, device: &Device<B>) -> BoolTensor<B> {
    fn bool_slice_assign(
        tensor: BoolTensor<Self>,
        ranges: &[core::ops::Range<usize>],
        value: BoolTensor<Self>,
    ) -> BoolTensor<Self> {
    fn bool_cat(tensors: Vec<BoolTensor<B>>, dim: usize) -> BoolTensor<B> {
    fn bool_equal(lhs: BoolTensor<B>, rhs: BoolTensor<B>) -> BoolTensor<B> {
    fn bool_not(tensor: BoolTensor<B>) -> BoolTensor<B> {
    fn bool_and(lhs: BoolTensor<B>, rhs: BoolTensor<B>) -> BoolTensor<B> {
    fn bool_or(lhs: BoolTensor<B>, rhs: BoolTensor<B>) -> BoolTensor<B> {
    fn bool_into_float(tensor: BoolTensor<B>) -> <Autodiff<B> as Backend>::FloatTensorPrimitive {
    fn bool_swap_dims(
        tensor: <Autodiff<B> as Backend>::BoolTensorPrimitive,
        dim1: usize,
        dim2: usize,
    ) -> <Autodiff<B> as Backend>::BoolTensorPrimitive {
    fn bool_permute(tensor: BoolTensor<Self>, axes: &[usize]) -> BoolTensor<Self> {
    fn bool_flip(tensor: BoolTensor<B>, axes: &[usize]) -> BoolTensor<B> {
    async fn bool_argwhere(tensor: BoolTensor<B>) -> IntTensor<B> {
    fn bool_expand(tensor: BoolTensor<B>, shape: Shape) -> BoolTensor<B> {
    fn bool_repeat_dim(tensor: BoolTensor<B>, dim: usize, times: usize) -> BoolTensor<B> {
impl<B: Backend, C: CheckpointStrategy> IntTensorOps<Self> for Autodiff<B, C> {
    fn int_from_data(data: TensorData, device: &Device<Self>) -> IntTensor<B> {
    async fn int_into_data(tensor: IntTensor<B>) -> TensorData {
    fn int_to_device(tensor: IntTensor<B>, device: &Device<Self>) -> IntTensor<B> {
    fn int_device(tensor: &IntTensor<B>) -> Device<Self> {
    fn int_reshape(tensor: IntTensor<B>, shape: Shape) -> IntTensor<B> {
    fn int_slice(tensor: IntTensor<B>, ranges: &[core::ops::Range<usize>]) -> IntTensor<B> {
    fn int_empty(shape: Shape, device: &<Autodiff<B> as Backend>::Device) -> IntTensor<B> {
    fn int_slice_assign(
        tensor: IntTensor<B>,
        ranges: &[core::ops::Range<usize>],
        value: IntTensor<B>,
    ) -> IntTensor<B> {
    fn int_cat(tensors: Vec<IntTensor<B>>, dim: usize) -> IntTensor<B> {
    fn int_equal(lhs: IntTensor<B>, rhs: IntTensor<B>) -> BoolTensor<B> {
    fn int_equal_elem(lhs: IntTensor<B>, rhs: B::IntElem) -> BoolTensor<B> {
    fn int_add(lhs: IntTensor<B>, rhs: IntTensor<B>) -> IntTensor<B> {
    fn int_add_scalar(lhs: IntTensor<B>, rhs: B::IntElem) -> IntTensor<B> {
    fn int_clamp_min(tensor: IntTensor<B>, min: B::IntElem) -> IntTensor<B> {
    fn int_clamp_max(tensor: IntTensor<B>, max: B::IntElem) -> IntTensor<B> {
    fn int_clamp(tensor: IntTensor<B>, min: B::IntElem, max: B::IntElem) -> IntTensor<B> {
    fn int_sub(lhs: IntTensor<B>, rhs: IntTensor<B>) -> IntTensor<B> {
    fn int_sub_scalar(lhs: IntTensor<B>, rhs: B::IntElem) -> IntTensor<B> {
    fn int_mul(lhs: IntTensor<B>, rhs: IntTensor<B>) -> IntTensor<B> {
    fn int_mul_scalar(lhs: IntTensor<B>, rhs: B::IntElem) -> IntTensor<B> {
    fn int_div(lhs: IntTensor<B>, rhs: IntTensor<B>) -> IntTensor<B> {
    fn int_div_scalar(lhs: IntTensor<B>, rhs: B::IntElem) -> IntTensor<B> {
    fn int_remainder(lhs: IntTensor<B>, rhs: IntTensor<B>) -> IntTensor<B> {
    fn int_remainder_scalar(lhs: IntTensor<B>, rhs: B::IntElem) -> IntTensor<B> {
    fn int_neg(tensor: IntTensor<B>) -> IntTensor<B> {
    fn int_zeros(shape: Shape, device: &Device<Self>) -> IntTensor<B> {
    fn int_ones(shape: Shape, device: &Device<Self>) -> IntTensor<B> {
    fn int_full(shape: Shape, fill_value: B::IntElem, device: &Device<Self>) -> IntTensor<B> {
    fn int_sum(tensor: IntTensor<B>) -> IntTensor<B> {
    fn int_sum_dim(tensor: IntTensor<B>, dim: usize) -> IntTensor<B> {
    fn int_mean(tensor: IntTensor<B>) -> IntTensor<B> {
    fn int_mean_dim(tensor: IntTensor<B>, dim: usize) -> IntTensor<B> {
    fn int_repeat_dim(tensor: IntTensor<B>, dim: usize, times: usize) -> IntTensor<B> {
    fn int_greater(lhs: IntTensor<B>, rhs: IntTensor<B>) -> BoolTensor<B> {
    fn int_greater_elem(lhs: IntTensor<B>, rhs: B::IntElem) -> BoolTensor<B> {
    fn int_greater_equal(lhs: IntTensor<B>, rhs: IntTensor<B>) -> BoolTensor<B> {
    fn int_greater_equal_elem(lhs: IntTensor<B>, rhs: B::IntElem) -> BoolTensor<B> {
    fn int_lower(lhs: IntTensor<B>, rhs: IntTensor<B>) -> BoolTensor<B> {
    fn int_lower_elem(lhs: IntTensor<B>, rhs: B::IntElem) -> BoolTensor<B> {
    fn int_lower_equal(lhs: IntTensor<B>, rhs: IntTensor<B>) -> BoolTensor<B> {
    fn int_lower_equal_elem(lhs: IntTensor<B>, rhs: B::IntElem) -> BoolTensor<B> {
    fn int_gather(dim: usize, tensor: IntTensor<B>, indices: IntTensor<B>) -> IntTensor<B> {
    fn int_scatter(
        dim: usize,
        tensor: IntTensor<B>,
        indices: IntTensor<B>,
        value: IntTensor<B>,
    ) -> IntTensor<B> {
    fn int_select(tensor: IntTensor<B>, dim: usize, indices: IntTensor<B>) -> IntTensor<B> {
    fn int_select_assign(
        tensor: IntTensor<B>,
        dim: usize,
        indices: IntTensor<B>,
        value: IntTensor<B>,
    ) -> IntTensor<B> {
    fn int_mask_where(
        tensor: IntTensor<B>,
        mask: BoolTensor<B>,
        value: IntTensor<B>,
    ) -> <Autodiff<B> as Backend>::IntTensorPrimitive {
    fn int_mask_fill(
        tensor: IntTensor<B>,
        mask: BoolTensor<B>,
        value: B::IntElem,
    ) -> <Autodiff<B> as Backend>::IntTensorPrimitive {
    fn int_argmax(tensor: IntTensor<B>, dim: usize) -> IntTensor<B> {
    fn int_argmin(tensor: IntTensor<B>, dim: usize) -> IntTensor<B> {
    fn int_max(tensor: B::IntTensorPrimitive) -> B::IntTensorPrimitive {
    fn int_max_dim(tensor: B::IntTensorPrimitive, dim: usize) -> B::IntTensorPrimitive {
    fn int_max_dim_with_indices(
        tensor: B::IntTensorPrimitive,
        dim: usize,
    ) -> (B::IntTensorPrimitive, B::IntTensorPrimitive) {
    fn int_min(tensor: B::IntTensorPrimitive) -> B::IntTensorPrimitive {
    fn int_min_dim(tensor: B::IntTensorPrimitive, dim: usize) -> B::IntTensorPrimitive {
    fn int_min_dim_with_indices(
        tensor: B::IntTensorPrimitive,
        dim: usize,
    ) -> (B::IntTensorPrimitive, B::IntTensorPrimitive) {
    fn int_abs(tensor: B::IntTensorPrimitive) -> B::IntTensorPrimitive {
    fn int_into_float(
        tensor: <Autodiff<B> as Backend>::IntTensorPrimitive,
    ) -> <Autodiff<B> as Backend>::FloatTensorPrimitive {
    fn int_swap_dims(
        tensor: <Autodiff<B> as Backend>::IntTensorPrimitive,
        dim1: usize,
        dim2: usize,
    ) -> <Autodiff<B> as Backend>::IntTensorPrimitive {
    fn int_random(
        shape: Shape,
        distribution: Distribution,
        device: &Device<Self>,
    ) -> IntTensor<Self> {
    fn int_arange(range: core::ops::Range<i64>, device: &Device<Self>) -> IntTensor<Self> {
    fn int_permute(tensor: IntTensor<Self>, axes: &[usize]) -> IntTensor<Self> {
    fn int_flip(tensor: IntTensor<Self>, axes: &[usize]) -> IntTensor<Self> {
    fn int_sign(tensor: IntTensor<Self>) -> IntTensor<Self> {
    fn int_prod(tensor: IntTensor<Self>) -> IntTensor<Self> {
    fn int_prod_dim(tensor: IntTensor<Self>, dim: usize) -> IntTensor<Self> {
    fn int_expand(tensor: IntTensor<B>, shape: Shape) -> IntTensor<B> {
    fn int_sort(tensor: IntTensor<Self>, dim: usize, descending: bool) -> IntTensor<Self> {
    fn int_sort_with_indices(
        tensor: IntTensor<Self>,
        dim: usize,
        descending: bool,
    ) -> (IntTensor<Self>, IntTensor<Self>) {
    fn int_argsort(tensor: IntTensor<Self>, dim: usize, descending: bool) -> IntTensor<Self> {
    fn bitwise_and(lhs: IntTensor<Self>, rhs: IntTensor<Self>) -> IntTensor<Self> {
    fn bitwise_and_scalar(lhs: IntTensor<Self>, rhs: B::IntElem) -> IntTensor<Self> {
    fn bitwise_or(lhs: IntTensor<Self>, rhs: IntTensor<Self>) -> IntTensor<Self> {
    fn bitwise_or_scalar(lhs: IntTensor<Self>, rhs: B::IntElem) -> IntTensor<Self> {
    fn bitwise_xor(lhs: IntTensor<Self>, rhs: IntTensor<Self>) -> IntTensor<Self> {
    fn bitwise_xor_scalar(lhs: IntTensor<Self>, rhs: B::IntElem) -> IntTensor<Self> {
    fn bitwise_not(tensor: IntTensor<Self>) -> IntTensor<Self> {
    fn bitwise_left_shift(lhs: IntTensor<Self>, rhs: IntTensor<Self>) -> IntTensor<Self> {
    fn bitwise_left_shift_scalar(lhs: IntTensor<Self>, rhs: B::IntElem) -> IntTensor<Self> {
    fn bitwise_right_shift(lhs: IntTensor<Self>, rhs: IntTensor<Self>) -> IntTensor<Self> {
    fn bitwise_right_shift_scalar(lhs: IntTensor<Self>, rhs: B::IntElem) -> IntTensor<Self> {
pub(crate) struct MaxMinDim;
impl<B: Backend> Backward<B, 1> for MaxMinDim {
    type State = (B::IntTensorPrimitive, Shape);
    fn backward(
        self,
        ops: Ops<Self::State, 1>,
        grads: &mut Gradients,
        _checkpointer: &mut Checkpointer,
    ) {
impl<B: Backend, C: CheckpointStrategy> ModuleOps<Autodiff<B, C>> for Autodiff<B, C> {
    fn embedding(weights: AutodiffTensor<B>, indices: IntTensor<B>) -> AutodiffTensor<B> {
        struct Embedding;
        impl<B: Backend> Backward<B, 1> for Embedding {
            type State = (B::FloatTensorPrimitive, IntTensor<B>);
            fn backward(
                self,
                ops: Ops<Self::State, 1>,
                grads: &mut Gradients,
                _checkpointer: &mut Checkpointer,
            ) {
    fn embedding_backward(
        _weights: AutodiffTensor<B>,
        _output: AutodiffTensor<B>,
        _indices: IntTensor<B>,
    ) -> AutodiffTensor<B> {
    fn conv1d(
        x: AutodiffTensor<B>,
        weight: AutodiffTensor<B>,
        bias: Option<AutodiffTensor<B>>,
        options: ConvOptions<1>,
    ) -> AutodiffTensor<B> {
        struct Conv1DWithBias;
        struct Conv1DNoBias;
        impl<B: Backend> Backward<B, 3> for Conv1DWithBias {
            type State = (NodeID, NodeID, NodeID, ConvOptions<1>);
            fn backward(
                self,
                ops: Ops<Self::State, 3>,
                grads: &mut Gradients,
                checkpointer: &mut Checkpointer,
            ) {
        impl<B: Backend> Backward<B, 2> for Conv1DNoBias {
            type State = (NodeID, NodeID, ConvOptions<1>);
            fn backward(
                self,
                ops: Ops<Self::State, 2>,
                grads: &mut Gradients,
                checkpointer: &mut Checkpointer,
            ) {
    fn conv_transpose1d(
        x: AutodiffTensor<B>,
        weight: AutodiffTensor<B>,
        bias: Option<AutodiffTensor<B>>,
        options: ConvTransposeOptions<1>,
    ) -> AutodiffTensor<B> {
        struct ConvTranspose1DWithBias;
        struct ConvTranspose1DNoBias;
        impl<B: Backend> Backward<B, 3> for ConvTranspose1DWithBias {
            type State = (NodeID, NodeID, NodeID, ConvTransposeOptions<1>);
            fn backward(
                self,
                ops: Ops<Self::State, 3>,
                grads: &mut Gradients,
                checkpointer: &mut Checkpointer,
            ) {
        impl<B: Backend> Backward<B, 2> for ConvTranspose1DNoBias {
            type State = (NodeID, NodeID, ConvTransposeOptions<1>);
            fn backward(
                self,
                ops: Ops<Self::State, 2>,
                grads: &mut Gradients,
                checkpointer: &mut Checkpointer,
            ) {
    fn conv2d(
        x: AutodiffTensor<B>,
        weight: AutodiffTensor<B>,
        bias: Option<AutodiffTensor<B>>,
        options: ConvOptions<2>,
    ) -> AutodiffTensor<B> {
        struct Conv2DWithBias;
        struct Conv2DNoBias;
        impl<B: Backend> Backward<B, 3> for Conv2DWithBias {
            type State = (NodeID, NodeID, NodeID, ConvOptions<2>);
            fn backward(
                self,
                ops: Ops<Self::State, 3>,
                grads: &mut Gradients,
                checkpointer: &mut Checkpointer,
            ) {
        impl<B: Backend> Backward<B, 2> for Conv2DNoBias {
            type State = (NodeID, NodeID, ConvOptions<2>);
            fn backward(
                self,
                ops: Ops<Self::State, 2>,
                grads: &mut Gradients,
                checkpointer: &mut Checkpointer,
            ) {
    fn deform_conv2d(
        x: AutodiffTensor<B>,
        offset: AutodiffTensor<B>,
        weight: AutodiffTensor<B>,
        mask: Option<AutodiffTensor<B>>,
        bias: Option<AutodiffTensor<B>>,
        options: DeformConvOptions<2>,
    ) -> AutodiffTensor<B> {
        struct DeformConv2DWithMaskWithBias;
        struct DeformConv2DWithMaskNoBias;
        struct DeformConv2DNoMaskWithBias;
        struct DeformConv2DNoMaskNoBias;
        impl<B: Backend> Backward<B, 5> for DeformConv2DWithMaskWithBias {
            type State = (NodeID, NodeID, NodeID, NodeID, NodeID, DeformConvOptions<2>);
            fn backward(
                self,
                ops: Ops<Self::State, 5>,
                grads: &mut Gradients,
                checkpointer: &mut Checkpointer,
            ) {
        impl<B: Backend> Backward<B, 4> for DeformConv2DWithMaskNoBias {
            type State = (NodeID, NodeID, NodeID, NodeID, DeformConvOptions<2>);
            fn backward(
                self,
                ops: Ops<Self::State, 4>,
                grads: &mut Gradients,
                checkpointer: &mut Checkpointer,
            ) {
        impl<B: Backend> Backward<B, 4> for DeformConv2DNoMaskWithBias {
            type State = (NodeID, NodeID, NodeID, NodeID, DeformConvOptions<2>);
            fn backward(
                self,
                ops: Ops<Self::State, 4>,
                grads: &mut Gradients,
                checkpointer: &mut Checkpointer,
            ) {
        impl<B: Backend> Backward<B, 3> for DeformConv2DNoMaskNoBias {
            type State = (NodeID, NodeID, NodeID, DeformConvOptions<2>);
            fn backward(
                self,
                ops: Ops<Self::State, 3>,
                grads: &mut Gradients,
                checkpointer: &mut Checkpointer,
            ) {
    fn deform_conv2d_backward(
        _x: AutodiffTensor<B>,
        _offset: AutodiffTensor<B>,
        _weight: AutodiffTensor<B>,
        _mask: Option<AutodiffTensor<B>>,
        _bias: Option<AutodiffTensor<B>>,
        _output_grad: AutodiffTensor<B>,
        _options: DeformConvOptions<2>,
    ) -> DeformConv2dBackward<Self> {
    fn conv_transpose2d(
        x: AutodiffTensor<B>,
        weight: AutodiffTensor<B>,
        bias: Option<AutodiffTensor<B>>,
        options: ConvTransposeOptions<2>,
    ) -> AutodiffTensor<B> {
        struct ConvTranspose2DWithBias;
        struct ConvTranspose2DNoBias;
        impl<B: Backend> Backward<B, 3> for ConvTranspose2DWithBias {
            type State = (NodeID, NodeID, NodeID, ConvTransposeOptions<2>);
            fn backward(
                self,
                ops: Ops<Self::State, 3>,
                grads: &mut Gradients,
                checkpointer: &mut Checkpointer,
            ) {
        impl<B: Backend> Backward<B, 2> for ConvTranspose2DNoBias {
            type State = (NodeID, NodeID, ConvTransposeOptions<2>);
            fn backward(
                self,
                ops: Ops<Self::State, 2>,
                grads: &mut Gradients,
                checkpointer: &mut Checkpointer,
            ) {
    fn conv3d(
        x: AutodiffTensor<B>,
        weight: AutodiffTensor<B>,
        bias: Option<AutodiffTensor<B>>,
        options: ConvOptions<3>,
    ) -> AutodiffTensor<B> {
        struct Conv3DWithBias;
        struct Conv3DNoBias;
        impl<B: Backend> Backward<B, 3> for Conv3DWithBias {
            type State = (NodeID, NodeID, NodeID, ConvOptions<3>);
            fn backward(
                self,
                ops: Ops<Self::State, 3>,
                grads: &mut Gradients,
                checkpointer: &mut Checkpointer,
            ) {
        impl<B: Backend> Backward<B, 2> for Conv3DNoBias {
            type State = (NodeID, NodeID, ConvOptions<3>);
            fn backward(
                self,
                ops: Ops<Self::State, 2>,
                grads: &mut Gradients,
                checkpointer: &mut Checkpointer,
            ) {
    fn conv_transpose3d(
        x: AutodiffTensor<B>,
        weight: AutodiffTensor<B>,
        bias: Option<AutodiffTensor<B>>,
        options: ConvTransposeOptions<3>,
    ) -> AutodiffTensor<B> {
        struct ConvTranspose3DWithBias;
        struct ConvTranspose3DNoBias;
        impl<B: Backend> Backward<B, 3> for ConvTranspose3DWithBias {
            type State = (NodeID, NodeID, NodeID, ConvTransposeOptions<3>);
            fn backward(
                self,
                ops: Ops<Self::State, 3>,
                grads: &mut Gradients,
                checkpointer: &mut Checkpointer,
            ) {
        impl<B: Backend> Backward<B, 2> for ConvTranspose3DNoBias {
            type State = (NodeID, NodeID, ConvTransposeOptions<3>);
            fn backward(
                self,
                ops: Ops<Self::State, 2>,
                grads: &mut Gradients,
                checkpointer: &mut Checkpointer,
            ) {
// TODO: Support a custom unfold4d operation by overriding the default implementation.
//
// We don't override it now because the fold operation isn't available for the backward pass.
// This implies that when autodiff is enabled, custom unfold operations defined by backends
// won't be used. Instead, the conv2d operation with custom weights matrix will be used.
// Therefore, the conv2d backward pass will be used for the unfold4d backward pass.
//
// fn unfold4d(
//     x:AutodiffTensor<B>,
//     kernel_size: [usize; 2],
//     options: UnfoldOptions,
// ) -> AutodiffTensor<B> {
//     todo!()
// }
    fn avg_pool1d(
        x: AutodiffTensor<B>,
        kernel_size: usize,
        stride: usize,
        padding: usize,
        count_include_pad: bool,
    ) -> AutodiffTensor<B> {
        struct AvgPool1D;
        impl<B: Backend> Backward<B, 1> for AvgPool1D {
            type State = (NodeID, usize, usize, usize, bool);
            fn backward(
                self,
                ops: Ops<Self::State, 1>,
                grads: &mut Gradients,
                checkpointer: &mut Checkpointer,
            ) {
    fn avg_pool2d(
        x: AutodiffTensor<B>,
        kernel_size: [usize;
        struct AvgPool2D;
        impl<B: Backend> Backward<B, 1> for AvgPool2D {
            type State = (NodeID, [usize;
            fn backward(
                self,
                ops: Ops<Self::State, 1>,
                grads: &mut Gradients,
                checkpointer: &mut Checkpointer,
            ) {
    fn avg_pool2d_backward(
        _x: AutodiffTensor<B>,
        _grad: AutodiffTensor<B>,
        _kernel_size: [usize;
    fn max_pool1d(
        x: AutodiffTensor<B>,
        kernel_size: usize,
        stride: usize,
        padding: usize,
        dilation: usize,
    ) -> AutodiffTensor<B> {
    fn max_pool1d_with_indices(
        x: AutodiffTensor<B>,
        kernel_size: usize,
        stride: usize,
        padding: usize,
        dilation: usize,
    ) -> MaxPool1dWithIndices<Self> {
    fn max_pool1d_with_indices_backward(
        x: AutodiffTensor<B>,
        kernel_size: usize,
        stride: usize,
        padding: usize,
        dilation: usize,
        output_grad: AutodiffTensor<B>,
        indices: IntTensor<B>,
    ) -> MaxPool1dBackward<Self> {
    fn max_pool2d(
        x: AutodiffTensor<B>,
        kernel_size: [usize;
    fn max_pool2d_with_indices(
        x: AutodiffTensor<B>,
        kernel_size: [usize;
    fn max_pool2d_with_indices_backward(
        _x: AutodiffTensor<B>,
        _kernel_size: [usize;
    fn adaptive_avg_pool1d(x: AutodiffTensor<B>, output_size: usize) -> AutodiffTensor<B> {
        struct AdaptiveAvgPool1D;
        impl<B: Backend> Backward<B, 1> for AdaptiveAvgPool1D {
            type State = NodeID;
            fn backward(
                self,
                ops: Ops<Self::State, 1>,
                grads: &mut Gradients,
                checkpointer: &mut Checkpointer,
            ) {
    fn adaptive_avg_pool2d(x: AutodiffTensor<B>, output_size: [usize;
        struct AdaptiveAvgPool2D;
        impl<B: Backend> Backward<B, 1> for AdaptiveAvgPool2D {
            type State = NodeID;
            fn backward(
                self,
                ops: Ops<Self::State, 1>,
                grads: &mut Gradients,
                checkpointer: &mut Checkpointer,
            ) {
    fn adaptive_avg_pool2d_backward(
        _x: AutodiffTensor<B>,
        _grad: AutodiffTensor<B>,
    ) -> <Autodiff<B> as Backend>::FloatTensorPrimitive {
    fn interpolate(
        x: AutodiffTensor<B>,
        output_size: [usize;
        struct Interpolate;
        impl<B: Backend> Backward<B, 1> for Interpolate {
            type State = (NodeID, [usize;
            fn backward(
                self,
                ops: Ops<Self::State, 1>,
                grads: &mut Gradients,
                checkpointer: &mut Checkpointer,
            ) {
    fn interpolate_backward(
        _x: FloatTensor<Autodiff<B, C>>,
        _grad: FloatTensor<Autodiff<B, C>>,
        _output_size: [usize;
struct MaxPool1D;
impl<B: Backend> Backward<B, 1> for MaxPool1D {
    type State = (NodeID, IntTensor<B>, usize, usize, usize, usize);
    fn backward(
        self,
        ops: Ops<Self::State, 1>,
        grads: &mut Gradients,
        checkpointer: &mut Checkpointer,
    ) {
struct MaxPool2D;
impl<B: Backend> Backward<B, 1> for MaxPool2D {
    type State = (
        NodeID,
        IntTensor<B>,
        [usize;
    fn backward(
        self,
        ops: Ops<Self::State, 1>,
        grads: &mut Gradients,
        checkpointer: &mut Checkpointer,
    ) {
impl<B: Backend, C: CheckpointStrategy> QTensorOps<Self> for Autodiff<B, C> {
    fn q_from_data(_data: TensorData, _device: &Device<Self>) -> QuantizedTensor<Self> {
    fn quantize(
        _tensor: FloatTensor<Self>,
        _scheme: &QuantScheme,
        _qparams: QuantizationParametersPrimitive<Self>,
    ) -> QuantizedTensor<Self> {
// required for QAT
    fn quantize_dynamic(
        _tensor: FloatTensor<Self>,
        _scheme: &QuantScheme,
    ) -> QuantizedTensor<Self> {
    fn dequantize(_tensor: QuantizedTensor<Self>) -> FloatTensor<Self> {
    fn q_device(tensor: &QuantizedTensor<Self>) -> Device<Self> {
    fn q_to_device(
        _tensor: QuantizedTensor<Self>,
        _device: &Device<Self>,
    ) -> QuantizedTensor<Self> {
    fn q_reshape(tensor: QuantizedTensor<Self>, shape: Shape) -> QuantizedTensor<Self> {
    async fn q_into_data(tensor: QuantizedTensor<Self>) -> TensorData {
    fn q_swap_dims(
        _tensor: QuantizedTensor<Self>,
        _dim1: usize,
        _dim2: usize,
    ) -> QuantizedTensor<Self> {
    fn q_permute(_tensor: QuantizedTensor<Self>, _axes: &[usize]) -> QuantizedTensor<Self> {
    fn q_flip(_tensor: QuantizedTensor<Self>, _axes: &[usize]) -> QuantizedTensor<Self> {
    fn q_gather(
        _dim: usize,
        _tensor: QuantizedTensor<Self>,
        _indices: IntTensor<Self>,
    ) -> QuantizedTensor<Self> {
    fn q_select(
        _tensor: QuantizedTensor<Self>,
        _dim: usize,
        _indices: IntTensor<Self>,
    ) -> QuantizedTensor<Self> {
    fn q_slice(_tensor: QuantizedTensor<Self>, _ranges: &[Range<usize>]) -> QuantizedTensor<Self> {
    fn q_argmax(tensor: QuantizedTensor<Self>, dim: usize) -> IntTensor<Self> {
    fn q_argmin(tensor: QuantizedTensor<Self>, dim: usize) -> IntTensor<Self> {
    fn q_expand(_tensor: QuantizedTensor<Self>, _shape: Shape) -> QuantizedTensor<Self> {
pub(crate) struct SortDim;
impl<B: Backend> Backward<B, 1> for SortDim {
    type State = (B::IntTensorPrimitive, Shape);
    fn backward(
        self,
        ops: Ops<Self::State, 1>,
        grads: &mut Gradients,
        _checkpointer: &mut Checkpointer,
    ) {
// Unsqueeze op on primitive.
fn unsqueeze_like<B: Backend>(
    tensor: B::FloatTensorPrimitive,
    shape: Shape,
) -> B::FloatTensorPrimitive {
/*
    let mut dims = [1; D2];
    let num_ones = D2 - D;
    let shape = self.shape();
    dims[num_ones..(D + num_ones)].copy_from_slice(&shape.dims[..D]);
    let shape = Shape::new(dims);
    self.reshape(shape)
    */
impl<B: Backend, C: CheckpointStrategy> FloatTensorOps<Self> for Autodiff<B, C> {
    fn float_from_data(data: TensorData, device: &Device<Self>) -> FloatTensor<Self> {
    fn float_random(
        shape: Shape,
        distribution: burn_tensor::Distribution,
        device: &Device<Self>,
    ) -> FloatTensor<Self> {
    fn float_zeros(shape: Shape, device: &Device<Self>) -> FloatTensor<Self> {
    fn float_ones(shape: Shape, device: &Device<Self>) -> FloatTensor<Self> {
    async fn float_into_data(tensor: FloatTensor<Self>) -> TensorData {
    fn float_device(tensor: &FloatTensor<Self>) -> Device<Self> {
    fn float_to_device(tensor: FloatTensor<Self>, device: &Device<Self>) -> FloatTensor<Self> {
        struct ToDevice;
        impl<B: Backend> Backward<B, 1> for ToDevice {
            type State = B::Device;
            fn backward(
                self,
                ops: Ops<Self::State, 1>,
                grads: &mut Gradients,
                _checkpointer: &mut Checkpointer,
            ) {
    fn float_empty(shape: Shape, device: &Device<Self>) -> FloatTensor<Self> {
    fn float_add(lhs: FloatTensor<Self>, rhs: FloatTensor<Self>) -> FloatTensor<Self> {
        struct Add;
        impl<B: Backend> Backward<B, 2> for Add {
            type State = (Shape, Shape);
            fn backward(
                self,
                ops: Ops<Self::State, 2>,
                grads: &mut Gradients,
                _checkpointer: &mut Checkpointer,
            ) {
    fn float_add_scalar(lhs: FloatTensor<Self>, rhs: FloatElem<B>) -> FloatTensor<Self> {
        struct AddScalar;
        impl<B: Backend> Backward<B, 1> for AddScalar {
            type State = ();
            fn backward(
                self,
                ops: Ops<Self::State, 1>,
                grads: &mut Gradients,
                _checkpointer: &mut Checkpointer,
            ) {
    fn float_sub(lhs: FloatTensor<Self>, rhs: FloatTensor<Self>) -> FloatTensor<Self> {
        struct Sub;
        impl<B: Backend> Backward<B, 2> for Sub {
            type State = (Shape, Shape);
            fn backward(
                self,
                ops: Ops<Self::State, 2>,
                grads: &mut Gradients,
                _checkpointer: &mut Checkpointer,
            ) {
    fn float_sub_scalar(lhs: FloatTensor<Self>, rhs: FloatElem<B>) -> FloatTensor<Self> {
        struct SubScalar;
        impl<B: Backend> Backward<B, 1> for SubScalar {
            type State = ();
            fn backward(
                self,
                ops: Ops<Self::State, 1>,
                grads: &mut Gradients,
                _checkpointer: &mut Checkpointer,
            ) {
    fn float_mul(lhs: FloatTensor<Self>, rhs: FloatTensor<Self>) -> FloatTensor<Self> {
        struct Mul;
        impl<B: Backend> Backward<B, 2> for Mul {
            type State = (Option<NodeID>, Option<NodeID>, BinaryOpsBroadcast);
            fn backward(
                self,
                ops: Ops<Self::State, 2>,
                grads: &mut Gradients,
                checkpointer: &mut Checkpointer,
            ) {
    fn float_mul_scalar(lhs: FloatTensor<Self>, rhs: FloatElem<B>) -> FloatTensor<Self> {
        struct MulScalar;
        impl<B: Backend> Backward<B, 1> for MulScalar {
            type State = FloatElem<B>;
            fn backward(
                self,
                ops: Ops<Self::State, 1>,
                grads: &mut Gradients,
                _checkpointer: &mut Checkpointer,
            ) {
    fn float_div(lhs: FloatTensor<Self>, rhs: FloatTensor<Self>) -> FloatTensor<Self> {
        struct Div;
        impl<B: Backend> Backward<B, 2> for Div {
            type State = (Option<NodeID>, Option<NodeID>, BinaryOpsBroadcast);
            fn backward(
                self,
                ops: Ops<Self::State, 2>,
                grads: &mut Gradients,
                checkpointer: &mut Checkpointer,
            ) {
    fn float_div_scalar(lhs: FloatTensor<Self>, rhs: FloatElem<B>) -> FloatTensor<Self> {
        struct DivScalar;
        impl<B: Backend> Backward<B, 1> for DivScalar {
            type State = FloatElem<B>;
            fn backward(
                self,
                ops: Ops<Self::State, 1>,
                grads: &mut Gradients,
                _checkpointer: &mut Checkpointer,
            ) {
    fn float_remainder(lhs: FloatTensor<Self>, rhs: FloatTensor<Self>) -> FloatTensor<Self> {
        struct Rem;
        impl<B: Backend> Backward<B, 2> for Rem {
            type State = (Option<NodeID>, Option<NodeID>, BinaryOpsBroadcast);
            fn backward(
                self,
                ops: Ops<Self::State, 2>,
                grads: &mut Gradients,
                checkpointer: &mut Checkpointer,
            ) {
// remainder(x, y) = x - floor(x / y) * y
// partial(x - floor(x / y) * y, x) = 1
// partial(x - floor(x / y) * y, y) = - floor(x / y)
    fn float_remainder_scalar(lhs: FloatTensor<Self>, rhs: FloatElem<B>) -> FloatTensor<Self> {
        struct RemainderScalar;
        impl<B: Backend> Backward<B, 1> for RemainderScalar {
            type State = ();
            fn backward(
                self,
                ops: Ops<Self::State, 1>,
                grads: &mut Gradients,
                _checkpointer: &mut Checkpointer,
            ) {
    fn float_matmul(lhs: FloatTensor<Self>, rhs: FloatTensor<Self>) -> FloatTensor<Self> {
        struct Matmul;
        impl<B: Backend> Backward<B, 2> for Matmul {
            type State = (Option<NodeID>, Option<NodeID>, BinaryOpsBroadcast);
            fn backward(
                self,
                ops: Ops<Self::State, 2>,
                grads: &mut Gradients,
                checkpointer: &mut Checkpointer,
            ) {
    fn float_neg(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
        struct Neg;
        impl<B: Backend> Backward<B, 1> for Neg {
            type State = ();
            fn backward(
                self,
                ops: Ops<Self::State, 1>,
                grads: &mut Gradients,
                _checkpointer: &mut Checkpointer,
            ) {
    fn float_recip(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
        struct Recip;
        impl<B: Backend> Backward<B, 1> for Recip {
            type State = NodeID;
            fn backward(
                self,
                ops: Ops<Self::State, 1>,
                grads: &mut Gradients,
                checkpointer: &mut Checkpointer,
            ) {
    fn float_swap_dims(tensor: FloatTensor<Self>, dim1: usize, dim2: usize) -> FloatTensor<Self> {
        struct SwapDim;
        struct RetroSwapDims<B: Backend> {
        impl<B: Backend> RetroForward for RetroSwapDims<B> {
            fn forward(&self, states: &mut BackwardStates, out_node: NodeID) {
        impl<B: Backend> Backward<B, 1> for SwapDim {
            type State = (usize, usize);
            fn backward(
                self,
                ops: Ops<Self::State, 1>,
                grads: &mut Gradients,
                _checkpointer: &mut Checkpointer,
            ) {
    fn float_permute(tensor: FloatTensor<Self>, axes: &[usize]) -> FloatTensor<Self> {
        struct PermuteDim;
        struct RetroPermuteDims<B: Backend> {
        impl<B: Backend> RetroForward for RetroPermuteDims<B> {
            fn forward(&self, states: &mut BackwardStates, out_node: NodeID) {
        impl<B: Backend> Backward<B, 1> for PermuteDim {
            type State = Vec<usize>;
            fn backward(
                self,
                ops: Ops<Self::State, 1>,
                grads: &mut Gradients,
                _checkpointer: &mut Checkpointer,
            ) {
    fn float_flip(tensor: FloatTensor<Self>, axes: &[usize]) -> FloatTensor<Self> {
        struct FlipDim;
        struct RetroFlipDims<B: Backend> {
        impl<B: Backend> RetroForward for RetroFlipDims<B> {
            fn forward(&self, states: &mut BackwardStates, out_node: NodeID) {
        impl<B: Backend> Backward<B, 1> for FlipDim {
            type State = Vec<usize>;
            fn backward(
                self,
                ops: Ops<Self::State, 1>,
                grads: &mut Gradients,
                _checkpointer: &mut Checkpointer,
            ) {
    fn float_reshape(tensor: FloatTensor<Self>, shape: Shape) -> FloatTensor<Self> {
        struct ReshapeDim;
        struct RetroReshape<B: Backend> {
        impl<B: Backend> RetroForward for RetroReshape<B> {
            fn forward(&self, states: &mut BackwardStates, out_node: NodeID) {
        impl<B: Backend> Backward<B, 1> for ReshapeDim {
            type State = (Shape, Shape);
            fn backward(
                self,
                ops: Ops<Self::State, 1>,
                grads: &mut Gradients,
                _checkpointer: &mut Checkpointer,
            ) {
    fn float_gather(
        dim: usize,
        tensor: FloatTensor<Self>,
        indices: IntTensor<B>,
    ) -> FloatTensor<Self> {
        struct Gather;
        impl<B: Backend> Backward<B, 1> for Gather {
            type State = (usize, IntTensor<B>, Shape, B::Device);
            fn backward(
                self,
                ops: Ops<Self::State, 1>,
                grads: &mut Gradients,
                _checkpointer: &mut Checkpointer,
            ) {
    fn float_scatter(
        dim: usize,
        tensor: FloatTensor<Self>,
        indices: IntTensor<B>,
        value: FloatTensor<Self>,
    ) -> FloatTensor<Self> {
        struct Scatter;
        impl<B: Backend> Backward<B, 2> for Scatter {
            type State = (usize, IntTensor<B>, Shape, Shape, B::Device);
            fn backward(
                self,
                ops: Ops<Self::State, 2>,
                grads: &mut Gradients,
                _checkpointer: &mut Checkpointer,
            ) {
    fn float_select(
        tensor: FloatTensor<Self>,
        dim: usize,
        indices: IntTensor<B>,
    ) -> FloatTensor<Self> {
        struct Select;
        struct RetroSelect<B: Backend> {
        impl<B: Backend> RetroForward for RetroSelect<B> {
            fn forward(&self, states: &mut BackwardStates, out_node: NodeID) {
        impl<B: Backend> Backward<B, 1> for Select {
            type State = (usize, IntTensor<B>, Shape, B::Device);
            fn backward(
                self,
                ops: Ops<Self::State, 1>,
                grads: &mut Gradients,
                _checkpointer: &mut Checkpointer,
            ) {
    fn float_select_assign(
        tensor: FloatTensor<Self>,
        dim: usize,
        indices: IntTensor<B>,
        value: FloatTensor<Self>,
    ) -> FloatTensor<Self> {
        struct IndexSelectDimAssign;
        struct RetroSelectAssign<B: Backend> {
        impl<B: Backend> RetroForward for RetroSelectAssign<B> {
            fn forward(&self, states: &mut BackwardStates, out_node: NodeID) {
        impl<B: Backend> Backward<B, 2> for IndexSelectDimAssign {
            type State = (usize, IntTensor<B>);
            fn backward(
                self,
                ops: Ops<Self::State, 2>,
                grads: &mut Gradients,
                _checkpointer: &mut Checkpointer,
            ) {
    fn float_slice(
        tensor: FloatTensor<Self>,
        ranges: &[core::ops::Range<usize>],
    ) -> FloatTensor<Self> {
        struct Index;
        struct RetroSlice<B: Backend> {
        impl<B: Backend> RetroForward for RetroSlice<B> {
            fn forward(&self, states: &mut BackwardStates, out_node: NodeID) {
        impl<B: Backend> Backward<B, 1> for Index {
            type State = (Vec<core::ops::Range<usize>>, Shape, B::Device);
            fn backward(
                self,
                ops: Ops<Self::State, 1>,
                grads: &mut Gradients,
                _checkpointer: &mut Checkpointer,
            ) {
    fn float_slice_assign(
        tensor: FloatTensor<Self>,
        ranges: &[core::ops::Range<usize>],
        value: FloatTensor<Self>,
    ) -> FloatTensor<Self> {
        struct SliceAssign;
        struct RetroSliceAssign<B: Backend> {
        impl<B: Backend> RetroForward for RetroSliceAssign<B> {
            fn forward(&self, states: &mut BackwardStates, out_node: NodeID) {
        impl<B: Backend> Backward<B, 2> for SliceAssign {
            type State = (Vec<core::ops::Range<usize>>, Shape, B::Device);
            fn backward(
                self,
                ops: Ops<Self::State, 2>,
                grads: &mut Gradients,
                _checkpointer: &mut Checkpointer,
            ) {
    fn float_mask_where(
        tensor: FloatTensor<Self>,
        mask: BoolTensor<Self>,
        source: FloatTensor<Self>,
    ) -> FloatTensor<Self> {
        struct MaskWhere;
        impl<B: Backend> Backward<B, 2> for MaskWhere {
            type State = (BoolTensor<B>, Shape, Shape, B::Device);
            fn backward(
                self,
                ops: Ops<Self::State, 2>,
                grads: &mut Gradients,
                _checkpointer: &mut Checkpointer,
            ) {
    fn float_mask_fill(
        tensor: FloatTensor<Self>,
        mask: BoolTensor<B>,
        value: FloatElem<B>,
    ) -> FloatTensor<Self> {
        struct MaskFill;
        impl<B: Backend> Backward<B, 1> for MaskFill {
            type State = BoolTensor<B>;
            fn backward(
                self,
                ops: Ops<Self::State, 1>,
                grads: &mut Gradients,
                _checkpointer: &mut Checkpointer,
            ) {
    fn float_equal(lhs: FloatTensor<Self>, rhs: FloatTensor<Self>) -> BoolTensor<B> {
    fn float_equal_elem(lhs: FloatTensor<Self>, rhs: FloatElem<B>) -> BoolTensor<B> {
    fn float_greater(lhs: FloatTensor<Self>, rhs: FloatTensor<Self>) -> BoolTensor<B> {
    fn float_greater_elem(lhs: FloatTensor<Self>, rhs: FloatElem<B>) -> BoolTensor<B> {
    fn float_greater_equal(lhs: FloatTensor<Self>, rhs: FloatTensor<Self>) -> BoolTensor<B> {
    fn float_greater_equal_elem(lhs: FloatTensor<Self>, rhs: FloatElem<B>) -> BoolTensor<B> {
    fn float_lower(lhs: FloatTensor<Self>, rhs: FloatTensor<Self>) -> BoolTensor<B> {
    fn float_lower_elem(lhs: FloatTensor<Self>, rhs: FloatElem<B>) -> BoolTensor<B> {
    fn float_lower_equal(lhs: FloatTensor<Self>, rhs: FloatTensor<Self>) -> BoolTensor<B> {
    fn float_lower_equal_elem(lhs: FloatTensor<Self>, rhs: FloatElem<B>) -> BoolTensor<B> {
    fn float_detach(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
// When we detach a tensor, we remove it from the graph, but we still want to keep the
// `require_grad` setting.
    fn float_set_require_grad(tensor: FloatTensor<Self>, require_grad: bool) -> FloatTensor<Self> {
    fn float_is_require_grad(tensor: &FloatTensor<Self>) -> bool {
    fn float_mean(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
        struct Mean;
        impl<B: Backend> Backward<B, 1> for Mean {
            type State = Shape;
            fn backward(
                self,
                ops: Ops<Self::State, 1>,
                grads: &mut Gradients,
                _checkpointer: &mut Checkpointer,
            ) {
    fn float_sum(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
        struct Sum;
        impl<B: Backend> Backward<B, 1> for Sum {
            type State = Shape;
            fn backward(
                self,
                ops: Ops<Self::State, 1>,
                grads: &mut Gradients,
                _checkpointer: &mut Checkpointer,
            ) {
    fn float_mean_dim(tensor: FloatTensor<Self>, dim: usize) -> FloatTensor<Self> {
        struct MeanDim;
        impl<B: Backend> Backward<B, 1> for MeanDim {
            type State = (Shape, usize);
            fn backward(
                self,
                ops: Ops<Self::State, 1>,
                grads: &mut Gradients,
                _checkpointer: &mut Checkpointer,
            ) {
    fn float_sum_dim(tensor: FloatTensor<Self>, dim: usize) -> FloatTensor<Self> {
        struct SumDim;
        impl<B: Backend> Backward<B, 1> for SumDim {
            type State = (Shape, usize);
            fn backward(
                self,
                ops: Ops<Self::State, 1>,
                grads: &mut Gradients,
                _checkpointer: &mut Checkpointer,
            ) {
    fn float_argmax(tensor: FloatTensor<Self>, dim: usize) -> IntTensor<B> {
    fn float_argmin(tensor: FloatTensor<Self>, dim: usize) -> IntTensor<B> {
    fn float_exp(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
        struct Exp;
        impl<B: Backend> Backward<B, 1> for Exp {
            type State = NodeID;
            fn backward(
                self,
                ops: Ops<Self::State, 1>,
                grads: &mut Gradients,
                checkpointer: &mut Checkpointer,
            ) {
    fn float_log(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
        struct Log;
        impl<B: Backend> Backward<B, 1> for Log {
            type State = NodeID;
            fn backward(
                self,
                ops: Ops<Self::State, 1>,
                grads: &mut Gradients,
                checkpointer: &mut Checkpointer,
            ) {
    fn float_log1p(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
        struct Log1P;
        impl<B: Backend> Backward<B, 1> for Log1P {
            type State = NodeID;
            fn backward(
                self,
                ops: Ops<Self::State, 1>,
                grads: &mut Gradients,
                checkpointer: &mut Checkpointer,
            ) {
    fn float_powf_scalar(tensor: FloatTensor<Self>, value: f32) -> FloatTensor<Self> {
        struct PowfScalar;
        struct RetroPowfScalar<B: Backend> {
        impl<B: Backend> RetroForward for RetroPowfScalar<B> {
            fn forward(&self, states: &mut BackwardStates, out_node: NodeID) {
        impl<B: Backend> Backward<B, 1> for PowfScalar {
            type State = (NodeID, f32);
            fn backward(
                self,
                ops: Ops<Self::State, 1>,
                grads: &mut Gradients,
                checkpointer: &mut Checkpointer,
            ) {
    fn float_sqrt(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
        struct Sqrt;
        impl<B: Backend> Backward<B, 1> for Sqrt {
            type State = NodeID;
            fn backward(
                self,
                ops: Ops<Self::State, 1>,
                grads: &mut Gradients,
                checkpointer: &mut Checkpointer,
            ) {
    fn float_abs(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
        struct Abs;
        impl<B: Backend> Backward<B, 1> for Abs {
            type State = NodeID;
            fn backward(
                self,
                ops: Ops<Self::State, 1>,
                grads: &mut Gradients,
                checkpointer: &mut Checkpointer,
            ) {
    fn float_cos(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
        struct Cos;
        impl<B: Backend> Backward<B, 1> for Cos {
            type State = NodeID;
            fn backward(
                self,
                ops: Ops<Self::State, 1>,
                grads: &mut Gradients,
                checkpointer: &mut Checkpointer,
            ) {
    fn float_sin(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
        struct Sin;
        impl<B: Backend> Backward<B, 1> for Sin {
            type State = NodeID;
            fn backward(
                self,
                ops: Ops<Self::State, 1>,
                grads: &mut Gradients,
                checkpointer: &mut Checkpointer,
            ) {
    fn float_tanh(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
        struct Tanh;
        impl<B: Backend> Backward<B, 1> for Tanh {
            type State = NodeID;
            fn backward(
                self,
                ops: Ops<Self::State, 1>,
                grads: &mut Gradients,
                checkpointer: &mut Checkpointer,
            ) {
    fn float_round(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
        struct Round;
        impl<B: Backend> Backward<B, 1> for Round {
            type State = (Shape, B::Device);
            fn backward(
                self,
                ops: Ops<Self::State, 1>,
                grads: &mut Gradients,
                _checkpointer: &mut Checkpointer,
            ) {
    fn float_floor(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
        struct Floor;
        impl<B: Backend> Backward<B, 1> for Floor {
            type State = (Shape, B::Device);
            fn backward(
                self,
                ops: Ops<Self::State, 1>,
                grads: &mut Gradients,
                _checkpointer: &mut Checkpointer,
            ) {
    fn float_ceil(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
        struct Ceil;
        impl<B: Backend> Backward<B, 1> for Ceil {
            type State = (Shape, B::Device);
            fn backward(
                self,
                ops: Ops<Self::State, 1>,
                grads: &mut Gradients,
                _checkpointer: &mut Checkpointer,
            ) {
    fn float_erf(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
        struct Erf;
        impl<B: Backend> Backward<B, 1> for Erf {
            type State = NodeID;
            fn backward(
                self,
                ops: Ops<Self::State, 1>,
                grads: &mut Gradients,
                checkpointer: &mut Checkpointer,
            ) {
    fn float_cat(tensors: Vec<FloatTensor<Self>>, dim: usize) -> FloatTensor<Self> {
        struct CatStep<B: Backend> {
// The dimension of each tensor along the dim dimension.
// This indicates the number of dimension concatenated for each tensor.
        impl<B: Backend> Step for CatStep<B> {
            fn step(self: Box<Self>, grads: &mut Gradients, _checkpointer: &mut Checkpointer) {
            fn node(&self) -> NodeID {
            fn parents(&self) -> Vec<NodeID> {
            fn depth(&self) -> usize {
// For simplicity, this operation does not checkpoint anything
    fn float_max_dim(tensor: FloatTensor<Self>, dim: usize) -> FloatTensor<Self> {
    fn float_max_dim_with_indices(
        tensor: FloatTensor<Self>,
        dim: usize,
    ) -> (FloatTensor<Self>, IntTensor<B>) {
    fn float_min_dim(tensor: FloatTensor<Self>, dim: usize) -> FloatTensor<Self> {
    fn float_min_dim_with_indices(
        tensor: FloatTensor<Self>,
        dim: usize,
    ) -> (FloatTensor<Self>, IntTensor<B>) {
    fn float_into_int(tensor: FloatTensor<Self>) -> <Autodiff<B> as Backend>::IntTensorPrimitive {
    fn float_powf(lhs: FloatTensor<Self>, rhs: FloatTensor<Self>) -> FloatTensor<Self> {
        struct PowF;
        impl<B: Backend> Backward<B, 2> for PowF {
            type State = (NodeID, NodeID, BinaryOpsBroadcast);
            fn backward(
                self,
                ops: Ops<Self::State, 2>,
                grads: &mut Gradients,
                checkpointer: &mut Checkpointer,
            ) {
// Both lhs and rhs are needed for both lhs and rhs gradients, but we clone them
// the number of times required by the parents specification.
//rhs*(lhs.val**(rhs-1))*grad
//lhs**rhs * ln(lhs) * grad
    fn float_sign(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
        struct Sign;
        impl<B: Backend> Backward<B, 1> for Sign {
            type State = ();
            fn backward(
                self,
                ops: Ops<Self::State, 1>,
                grads: &mut Gradients,
                _checkpointer: &mut Checkpointer,
            ) {
// Always return 0 because the derivative of the sign function
// does not contribute to gradient updates in a meaningful way.
    fn float_expand(tensor: FloatTensor<Self>, shape: Shape) -> FloatTensor<Self> {
// D1: tensor, D2: shape
        struct ExpandDim;
        struct RetroExpand<B: Backend> {
        impl<B: Backend> RetroForward for RetroExpand<B> {
            fn forward(&self, states: &mut BackwardStates, out_node: NodeID) {
        impl<B: Backend> Backward<B, 1> for ExpandDim {
            type State = (Shape, Shape);
            fn backward(
                self,
                ops: Ops<Self::State, 1>,
                grads: &mut Gradients,
                _checkpointer: &mut Checkpointer,
            ) {
    fn float_sort(tensor: FloatTensor<Self>, dim: usize, descending: bool) -> FloatTensor<Self> {
    fn float_sort_with_indices(
        tensor: FloatTensor<Self>,
        dim: usize,
        descending: bool,
    ) -> (FloatTensor<Self>, IntTensor<B>) {
    fn float_argsort(tensor: FloatTensor<Self>, dim: usize, descending: bool) -> IntTensor<B> {
    fn float_repeat_dim(tensor: FloatTensor<Self>, dim: usize, times: usize) -> FloatTensor<Self> {
        struct Repeat;
        struct RetroRepeat<B: Backend> {
        impl<B: Backend> RetroForward for RetroRepeat<B> {
            fn forward(&self, states: &mut BackwardStates, out_node: NodeID) {
        impl<B: Backend> Backward<B, 1> for Repeat {
            type State = (usize, usize);
            fn backward(
                self,
                ops: Ops<Self::State, 1>,
                grads: &mut Gradients,
                _checkpointer: &mut Checkpointer,
            ) {
// shape [..., orig_dim_size, times, ...]
// sum over repeat times
    fn float_cast(tensor: FloatTensor<Self>, dtype: burn_tensor::FloatDType) -> FloatTensor<Self> {
// TODO: Implement float_prod and float_sum
// https://github.com/tracel-ai/burn/issues/1458
enum BinaryOpsBroadcast {
impl BinaryOpsBroadcast {
    fn new<B: Backend>(lhs: &B::FloatTensorPrimitive, rhs: &B::FloatTensorPrimitive) -> Self {
    fn backward_lhs<B: Backend>(&self, grad: B::FloatTensorPrimitive) -> B::FloatTensorPrimitive {
    fn backward_rhs<B: Backend>(&self, grad: B::FloatTensorPrimitive) -> B::FloatTensorPrimitive {
impl<B: Backend, C: CheckpointStrategy> TransactionOps<Self> for Autodiff<B, C> {
    async fn tr_execute(
        transaction: TransactionPrimitive<Self>,
    ) -> burn_tensor::ops::TransactionPrimitiveResult {
/// Client used to communicate with the autodiff server.
pub trait AutodiffClient: Send + Clone {
/// Register a new step.
    fn register(&self, node_id: NodeRefCount, step: StepBoxed, actions: CheckpointerBuilder);
/// Call backpropagation from the given tensor.
    fn backward<B: Backend>(&self, tensor: AutodiffTensor<B>) -> Gradients;
/// Client implementation in used.
pub type AutodiffClientImpl = super::mspc::ChannelClient;
/// Client implementation in used.
pub type AutodiffClientImpl = super::mutex::MutexClient;
pub struct GraphMemoryManagement {
enum NodeMemoryStatus {
impl GraphMemoryManagement {
/// Register a new node with its parent.
    pub fn register(&mut self, node: NodeRefCount, parents: Vec<NodeID>) {
/// Free the node from the state.
    pub fn consume_node(&mut self, node_id: NodeID) {
/// Free all nodes whose backward call has become impossible
///
/// This function goes into three steps, which must happen for all leaves
/// before going into the next step. Then it deletes what can be safely deleted
    pub(crate) fn free_unavailable_nodes(&mut self, mut on_free_graph: impl FnMut(&NodeID)) {
// When consuming nodes with a backward pass, some other backward passes become
// unavailable because some of their parents have been consumed. They are
// identified here.
// Among the available nodes that remain, some may be useless if no
// available node with a tensor reference exist in their descendance.
// But some may seem useless from some leaf but be useful from another one,
// hence the need to iterate on all leaves.
// New leaves are the roots of a useful backward sub-tree.
// Deletables are everything not marked as useful.
// Replace leaves by the new ones and delete everything not useful anymore
    fn clear_unused_roots(&mut self, to_delete: &mut Vec<NodeID>) {
// Check if parents are either empty or absent from self.nodes
    fn unavailable_propagation(&mut self, node_id: NodeID) -> NodeMemoryStatus {
// If already visited
// If node exists and any of its parents is unavailable, it is unavailable as well
// If node exists but the parents vec is empty, it is a tensor that never had parents;
//  the status remains unknown
// If node does not exist, it was
// deleted, so this and all its descendants are unavailable
    fn useful_propagation(&mut self, leaves: HashSet<NodeID>) {
// Accumulate visited nodes
// Queue of nodes to visit
// Utilitary function to iterate over a node's parents
// Pop a node id, greedily looking at tag_useful ones first
// There are no nodes in the queues anymore
// The node can be explored, as long as it's not already tagged useful
    fn identify_leaves_and_deletables(
        &self,
        leaf_id: NodeID,
        new_leaves: &mut HashSet<NodeID>,
        to_delete: &mut Vec<NodeID>,
    ) {
    fn is_referenced(&self, node_id: NodeID) -> bool {
/// Wrapper over hash set for fast popping of any node
struct PopNodeSet {
impl PopNodeSet {
    fn pop(&mut self) -> Option<NodeID> {
    fn contains(&self, node_id: &NodeID) -> bool {
    fn insert(&mut self, node_id: NodeID) {
static INSTANCE: spin::Lazy<ChannelClient> = spin::Lazy::new(ChannelClient::init);
pub struct ChannelClient {
enum Message {
impl ChannelClient {
    pub(crate) fn new() -> Self {
    fn init() -> Self {
impl AutodiffClient for ChannelClient {
    fn register(&self, node_id: NodeRefCount, step: StepBoxed, actions: CheckpointerBuilder) {
    fn backward<B: Backend>(&self, root: AutodiffTensor<B>) -> Gradients {
pub struct MutexClient;
impl core::fmt::Debug for MutexClient {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
static SERVER: spin::Mutex<Option<AutodiffServer>> = spin::Mutex::new(None);
impl AutodiffClient for MutexClient {
    fn register(&self, node_id: NodeRefCount, step: StepBoxed, actions: CheckpointerBuilder) {
    fn backward<B: Backend>(&self, root: AutodiffTensor<B>) -> Gradients {
pub struct AutodiffServer {
impl AutodiffServer {
    pub fn register(&mut self, rc: NodeRefCount, step: StepBoxed, actions: CheckpointerBuilder) {
    pub fn backward(&mut self, grads: Gradients, node_id: NodeID) -> Gradients {
// Cleanup
    fn build_tape(
        &mut self,
        node: NodeID,
        node_step: StepBoxed,
        mut builder: CheckpointerBuilder,
    ) -> (Vec<Vec<StepBoxed>>, Checkpointer) {
    fn execute_steps(
        tape: Vec<Vec<StepBoxed>>,
        mut grads: Gradients,
        mut checkpointer: Checkpointer,
    ) -> Gradients {
// For checkpointing tests
pub struct AutodiffTensor<B: Backend> {
impl<B: Backend> TensorMetadata for AutodiffTensor<B> {
    fn dtype(&self) -> burn_tensor::DType {
    fn shape(&self) -> burn_tensor::Shape {
pub type NodeRefCount = Arc<NodeID>;
pub(crate) struct RootStep {
impl Step for RootStep {
    fn step(self: Box<Self>, _grads: &mut Gradients, _checkpointer: &mut Checkpointer) {
// Nothing to do
    fn node(&self) -> NodeID {
    fn parents(&self) -> Vec<NodeID> {
    fn depth(&self) -> usize {
impl<B: Backend> AutodiffTensor<B> {
/// Create a new leaf tensor.
    pub fn new(primitive: B::FloatTensorPrimitive) -> Self {
    pub fn is_tracked(&self) -> bool {
/// Mark the tensor as requiring gradients.
///
/// # Panics
///
/// It panics if the tensor is not a leaf.
    pub fn require_grad(mut self) -> Self {
/// Create a tensor from parent infos.
    pub fn from_parents(
        primitive: B::FloatTensorPrimitive,
        parent_nodes: &[NodeRef],
        requirement: Requirement,
        computing_properties: ComputingProperty,
    ) -> Self {
/// Register a step into a graph for that tensor.
///
/// # Warning
///
/// This should be called only once per tensor.
    pub fn register_step<S: Step + 'static>(
        self,
        step_that_created_the_tensor: S,
        actions: CheckpointerBuilder,
    ) -> Self {
    pub fn into_primitive(self) -> B::FloatTensorPrimitive {
    fn should_diff_abs() {
    fn should_diff_abs_no_nans() {
    fn test_avg_pool1d_simple() {
    struct AdaptiveAvgPool1dTestCase {
    impl AdaptiveAvgPool1dTestCase {
        fn assert_output(self, x_grad: TestTensor<3>) {
    fn test_avg_pool2d_simple() {
    struct AdaptiveAvgPool2dTestCase {
    impl AdaptiveAvgPool2dTestCase {
        fn assert_output(self, x_grad: TestTensor<4>) {
    fn should_diff_add() {
    fn should_diff_add_scalar() {
    fn test_add_complex_1() {
    fn should_diff_mean() {
    fn should_diff_sum_1() {
    fn should_diff_sum_2() {
    fn should_diff_mean_dim() {
    fn should_diff_sum_dim() {
    fn test_avg_pool1d_simple() {
    fn test_avg_pool1d_complex() {
    fn test_avg_pool1d_complex_dont_count_pad() {
    struct AvgPool1dTestCase {
    impl AvgPool1dTestCase {
        fn assert_output(self, x_grad: TestTensor<3>) {
    fn test_avg_pool2d_simple() {
    fn test_avg_pool2d_complex() {
    fn test_avg_pool2d_complex_dont_include_pad() {
    struct AvgPool2dTestCase {
    impl AvgPool2dTestCase {
        fn assert_output(self, x_grad: TestTensor<4>) {
    fn test_embedding_backward() {
    fn test_full_precision() {
    fn mul_broadcast() {
    fn div_broadcast() {
    fn sub_broadcast() {
    fn add_broadcast() {
    fn matmul_broadcast() {
    fn mask_where_broadcast() {
    fn test_ops_broadcast_backward<F>(func: F)
    where
        F: Fn(TestAutodiffTensor<3>, TestAutodiffTensor<3>) -> TestAutodiffTensor<3>,
    {
// Slice isn't a broadcastable operation, so it will fail when the previous backward pass
// of an operation that support broadcast doesn't support it during the backward pass.
// Will panic if broadcast isn't supported!
    type FT = FloatElem<TestBackend>;
    fn should_diff_cat() {
    fn should_diff_cat_more_than_1_dim() {
// Concat a tensor [2, 2] with another tensor [3, 2] along dim 0.
// The resulting tensor should be [5, 2]
    fn should_diff_ceil() {
    fn test_autodiff_checkpoint_complicated_computation() {
    fn test_autodiff_checkpoint_with_missing_requirement() {
// does not require_grad
    fn test_autodiff_checkpoint_with_many_duplicates() {
    fn test_autodiff_checkpoint_with_long_chain_of_eager_memory_bound() {
    fn test_autodiff_checkpoint_half_sub_graph_not_tracked() {
    fn test_autodiff_checkpoint_very_complex() {
    fn assert_checkpoint<const D: usize>(tensor: TestAutodiffTensor<D>) {
// Assert is not explicit here, but the test can fail
// - when a tensor is actually required more than n_required, it won't be found and will panic
// - when a tensor is actually required less than n_required, the backward states map won't be
//   empty and will fail the assertion within the backward code, same for retro_forwards
// Does not save its state and does not need its parents
    fn memory_bound_eager<const D: usize>(
        tensor_a: TestAutodiffTensor<D>,
        tensor_b: TestAutodiffTensor<D>,
    ) -> TestAutodiffTensor<D> {
    fn memory_bound_eager_scalar<const D: usize>(
        tensor_a: TestAutodiffTensor<D>,
        b: f32,
    ) -> TestAutodiffTensor<D> {
// Saves its own state and does not need its parents
    fn compute_bound_eager<const D: usize>(
        tensor_a: TestAutodiffTensor<D>,
        tensor_b: TestAutodiffTensor<D>,
    ) -> TestAutodiffTensor<D> {
    fn compute_bound_eager_scalar<const D: usize>(
        tensor_a: TestAutodiffTensor<D>,
        b: f32,
    ) -> TestAutodiffTensor<D> {
// Does not save its state and needs its parents
    fn memory_bound_lazy<const D: usize>(
        tensor_a: TestAutodiffTensor<D>,
        tensor_b: TestAutodiffTensor<D>,
    ) -> TestAutodiffTensor<D> {
// Saves its own state and needs its parents
    fn compute_bound_lazy<const D: usize>(
        tensor_a: TestAutodiffTensor<D>,
        tensor_b: TestAutodiffTensor<D>,
    ) -> TestAutodiffTensor<D> {
    fn should_diff_full_complex_1() {
    fn should_diff_full_complex_2() {
    fn should_diff_full_complex_3() {
    type FT = FloatElem<TestBackend>;
    fn test_conv1d_basic() {
    fn test_conv1d_different_channels() {
    fn test_conv1d_with_padding() {
    fn test_conv1d_with_stride() {
    fn test_conv1d_dilation() {
    fn test_conv1d_groups() {
    struct Conv1dTestCase {
    struct Grads {
    impl Conv1dTestCase {
        fn assert_grads(self, expected_grads: Grads) {
// Assert
    type FT = FloatElem<TestBackend>;
    fn test_conv2d_basic() {
    fn test_conv2d_different_channels() {
    fn test_conv2d_different_kernel_size() {
    fn test_conv2d_different_padding() {
    fn test_conv2d_different_width() {
    fn test_conv2d_stride_2() {
    fn test_conv2d_different_stride() {
    fn test_conv2d_dilation_2() {
    fn test_conv2d_different_dilation() {
    fn test_conv2d_groups() {
    fn test_conv2d_groups_stride_2() {
    fn test_conv2d_groups_different_channels() {
    fn test_conv2d_complex() {
    struct Conv2dTestCase {
    struct Grads {
    impl Conv2dTestCase {
        fn assert_grads(self, expected_grads: Grads) {
// Assert
    type FT = FloatElem<TestBackend>;
    fn test_conv3d_basic() {
// TODO
    fn test_conv3d_complex() {
    struct Conv3dTestCase {
    struct Grads {
    impl Conv3dTestCase {
        fn assert_grads(self, expected_grads: Grads) {
// Assert
    type FT = FloatElem<TestBackend>;
    fn test_conv_transpose1d_basic() {
    fn test_conv_transpose1d_padding() {
    fn test_conv_transpose1d_stride() {
    fn test_conv_transpose1d_stride_padding_out() {
    fn test_conv_transpose1d_dilation() {
    fn test_conv_transpose1d_complex() {
    struct ConvTranspose1dTestCase {
    struct Grads {
    impl ConvTranspose1dTestCase {
        fn assert_grads(self, expected_grads: Grads) {
// Assert
    type FT = FloatElem<TestBackend>;
    fn test_conv_transpose2d_basic() {
    fn test_conv_transpose2d_padding() {
    fn test_conv_transpose2d_stride() {
    fn test_conv_transpose2d_stride_padding_out() {
    fn test_conv_transpose2d_dilation() {
    fn test_conv_transpose2d_channels() {
    fn test_conv_transpose2d_kernel_size() {
    fn test_conv_transpose2d_groups() {
    fn test_conv_transpose2d_complex_no_groups() {
    fn test_conv_transpose2d_complex_no_groups_2() {
    fn test_conv_transpose2d_complex_groups() {
    struct ConvTranspose2dTestCase {
    struct Grads {
    impl ConvTranspose2dTestCase {
        fn assert_grads(self, expected_grads: Grads) {
// Assert
    fn test_conv_transpose3d_basic() {
    fn test_conv_transpose3d_complex_groups() {
    struct ConvTranspose3dTestCase {
    struct Grads {
    impl ConvTranspose3dTestCase {
        fn assert_grads(self, expected_grads: Grads) {
// Assert
    fn should_diff_cos() {
// Metal has less precise trigonometric functions
    fn test_cross_entropy_loss_grad() {
    fn test_deform_conv2d_basic() {
    fn test_deform_conv2d_batched() {
    fn test_deform_conv2d_different_kernel_size() {
    fn test_deform_conv2d_different_padding() {
// => Position 788: 10.421875 != 10.0546875
//  diff (rel = +1.79e-2, abs = +3.67e-1), tol (rel = +1.00e-2, abs = +9.77e-4)
    struct Conv2dTestCase {
    struct Grads {
    impl Conv2dTestCase {
        fn assert_grads(self, expected_grads: Grads) {
// Assert
// Relative is set to 5%, which is much higher than typical numerical test tolerances.
// This is due to the complexity of the deformable convolution operation.
// Unlike regular conv2d, which samples from fixed integer grid positions,
// deformable conv2d samples input values at fractional offset locations (learned offsets).
// These non-integer positions require bilinear interpolation to estimate the input value.
// Gradients computed through all these floating-point operations can compound numerical differences.
    fn should_diff_div() {
    fn should_diff_div_scalar() {
    fn test_div_complex_1() {
    fn test_div_complex_2() {
    type FT = FloatElem<TestBackend>;
    fn should_diff_erf() {
    type FT = FloatElem<TestBackend>;
    fn should_diff_exp() {
    fn should_diff_expand() {
// Python code to generate the test case values
// import torch
// x1 = torch.tensor([4.0, 7.0, 2.0, 3.0], requires_grad=True)
// x2 = torch.tensor([2.0, 4.5, 7.0, 3.0], requires_grad=True)
// y = x1.expand(4, 4)
// z = (x2 * y).sum()
// z.backward()
// print("x1", x1.grad)
// print("x2", x2.grad)
// Use unsqueeze to make tensor_2 have the same shape as tensor_3
    fn should_diff_flip() {
// 1x2x2
// 1x2x3
// 1x2x2
// 1x2x3
    fn should_diff_floor() {
    fn test_gather_grad() {
    fn test_scatter_grad() {
    fn should_diff_gelu() {
    fn should_update_tensor_when_grad_replace() {
    fn should_diff_log() {
    type FT = FloatElem<TestBackend>;
    fn should_diff_log1p() {
    type FT = FloatElem<TestBackend>;
    fn should_diff_log_sigmoid() {
    type FT = FloatElem<TestBackend>;
    fn should_diff_mask_fill() {
    fn should_diff_mask_where() {
    fn should_diff_matmul() {
    fn test_matmul_complex_1() {
    fn test_matmul_complex_2() {
    type FT = FloatElem<TestBackend>;
    fn should_diff_max_dim() {
    fn should_diff_min_dim() {
    type FT = FloatElem<TestBackend>;
    fn test_max_pool1d_simple() {
// Asserts
    fn test_max_pool1d_with_dilation() {
// Asserts
    fn test_max_pool1d_complex() {
// Asserts
    fn test_max_pool1d_complex_with_padding() {
// Asserts
    type FT = FloatElem<TestBackend>;
    fn test_max_pool2d_simple_1() {
// Asserts
    fn test_max_pool2d_simple_2() {
// Asserts
    fn test_max_pool2d_with_dilation() {
// Asserts
    fn test_max_pool2d_complex() {
// Asserts
    fn test_mm_independent_trees() {
// First tree
// Second tree
    fn test_mm_crossover_trees_root_unavailable() {
// First tree
// Second tree
    fn test_mm_crossover_trees_with_referred_subtree() {
// First tree
// Second tree
    fn test_mm_three_crossover_trees_last_still_usable() {
// First tree
// Third tree
// Second tree (in between)
    fn test_mm_three_crossover_trees_middle_one_unavailable() {
// First tree
// Third tree
// Second tree (in between)
    fn test_mm_self_referencing_tree() {
// First tree
    fn test_mm_with_non_impacting_detach() {
    fn test_mm_with_missing_require_grad_after_cleanup() {
// Trivial backward, just to trigger cleanup
    fn test_mm_with_detach_after_cleanup() {
// Trivial backward, just to trigger cleanup
    fn test_mm_deletables_propagate_well() {
// We are testing that after backward on tensor_2, not only the leaf tensor_4 is deleted, but
// the intermediate tensor_3 as well
    fn test_mm_node_explored_once_can_still_be_tagged_as_useful_when_found_again_deeper() {
// The test has 50% chance of starting with leaf tensor_8 instead of tensor_4, which is not informative
// By repeating it many times it becomes almost impossible that it passes if it shouldn't
// tensor_2 should be tagged unknown through the leaf tensor_4, then useful through the leaf tensor_8
// which should happen after because tensor_2 is deeper from tensor_8 point of view and we're in breadth first search
// Avoid using paste dependency with no parameters
            type TestAutodiffBackend = burn_autodiff::Autodiff<TestBackend>;
            type TestAutodiffTensor<const D: usize> = burn_tensor::Tensor<TestAutodiffBackend, D>;
            pub type FloatType = <TestBackend as burn_tensor::backend::Backend>::FloatElem;
            pub type IntType = <TestBackend as burn_tensor::backend::Backend>::IntElem;
            pub type BoolType = <TestBackend as burn_tensor::backend::Backend>::BoolTensorPrimitive;
            type TestAutodiffBackend = burn_autodiff::Autodiff<TestBackend, burn_autodiff::checkpoint::strategy::BalancedCheckpointing>;
            type TestAutodiffTensor<const D: usize> = burn_tensor::Tensor<TestAutodiffBackend, D>;
            pub type FloatType = <TestBackend as burn_tensor::backend::Backend>::FloatElem;
            pub type IntType = <TestBackend as burn_tensor::backend::Backend>::IntElem;
            pub type BoolType = <TestBackend as burn_tensor::backend::Backend>::BoolTensorPrimitive;
            type TestAutodiffBackend = burn_autodiff::Autodiff<TestBackend, burn_autodiff::checkpoint::strategy::BalancedCheckpointing>;
            type TestAutodiffTensor<const D: usize> = burn_tensor::Tensor<TestAutodiffBackend, D>;
            pub type FloatType = <TestBackend as burn_tensor::backend::Backend>::FloatElem;
            pub type IntType = <TestBackend as burn_tensor::backend::Backend>::IntElem;
            pub type BoolType = <TestBackend as burn_tensor::backend::Backend>::BoolElem;
                    pub type TestBackend = TestBackend2<$float, IntType, BoolType>;
                    pub type TestAutodiffBackend = burn_autodiff::Autodiff<TestBackend, burn_autodiff::checkpoint::strategy::BalancedCheckpointing>;
                    pub type TestAutodiffTensor<const D: usize> = burn_tensor::Tensor<TestAutodiffBackend, D>;
                    pub type TestTensor<const D: usize> = TestTensor2<$float, IntType, BoolType, D>;
                    pub type TestTensorInt<const D: usize> = TestTensorInt2<$float, IntType, BoolType, D>;
                    pub type TestTensorBool<const D: usize> = TestTensorBool2<$float, IntType, BoolType, D>;
                    type FloatType = $float;
            type TestAutodiffBackend = burn_autodiff::Autodiff<TestBackend>;
            type TestAutodiffTensor<const D: usize> = burn_tensor::Tensor<TestAutodiffBackend, D>;
            pub type FloatType = <TestBackend as burn_tensor::backend::Backend>::FloatElem;
            pub type IntType = <TestBackend as burn_tensor::backend::Backend>::IntElem;
            pub type BoolType = <TestBackend as burn_tensor::backend::Backend>::BoolElem;
                    pub type TestBackend = TestBackend2<$float, IntType, BoolType>;
                    pub type TestAutodiffBackend = burn_autodiff::Autodiff<TestBackend>;
                    pub type TestAutodiffTensor<const D: usize> = burn_tensor::Tensor<TestAutodiffBackend, D>;
                    pub type TestTensor<const D: usize> = TestTensor2<$float, IntType, BoolType, D>;
                    pub type TestTensorInt<const D: usize> = TestTensorInt2<$float, IntType, BoolType, D>;
                    pub type TestTensorBool<const D: usize> = TestTensorBool2<$float, IntType, BoolType, D>;
                    type FloatType = $float;
// Behaviour
// Activation
// Modules
// Wgpu on MacOS currently doesn't support atomic compare exchange
// Tensor
    fn should_diff_mul() {
    fn should_diff_mul_scalar() {
    fn test_mul_complex_1() {
    fn should_behave_the_same_with_multithread() {
// Task 1
// Task 2
// Task 1
// Task 2
    type FT = FloatElem<TestBackend>;
    fn test_upsample_interpolation() {
    fn test_downsample_interpolation() {
    struct InterpolateTestCase {
    impl InterpolateTestCase {
        fn assert_output(self, x_grad: TestTensor<4>) {
    fn should_diff_neg() {
    fn should_diff_nonzero() {
// Multi-dimensional tensor indexing isn't really supported yet so the easiest way to do
// this is to flatten the mask and tensor to get proper indexing. Anyway the returned tensor would
// have dimensions different from the input, so this is somewhat equivalent.
// Vector dot product not supported (only 2D matmuls) so unsqueeze for test purposes
    type FT = FloatElem<TestBackend>;
    fn should_diff_permute() {
// 1x2x2
// 1x3x2
// 1x2x2
// 1x3x2
    type FT = FloatElem<TestBackend>;
    fn should_diff_powf_scalar() {
    fn should_diff_powf() {
    fn should_diff_powf_with_untracked_lhs() {
    fn should_diff_powf_with_untracked_rhs() {
    type FT = FloatElem<TestBackend>;
    fn should_diff_recip() {
    fn should_diff_relu() {
    type FT = FloatElem<TestBackend>;
    fn should_diff_remainder() {
    fn should_diff_repeat() {
    fn should_diff_repeat_multi_dim() {
    fn should_diff_reshape() {
    fn should_diff_round() {
    fn test_select_grad() {
    fn test_select_assign_grad() {
    fn test_select_assign_grad_different_shapes() {
    type FT = FloatElem<TestBackend>;
    fn should_diff_sigmoid() {
    fn small_neg_val_should_not_cause_grad_overflow() {
/// Example using the sign function with PyTorch:
// >>> import torch
// >>> # Create a tensor with requires_grad=True
// >>> x = torch.tensor([-2.0, -1.0, 0.0, 1.0, 2.0], requires_grad=True)
// >>> # Forward pass: Apply the sign function
// >>> y = torch.sign(x)
// >>> print("Forward pass:")
// Forward pass:
// >>> print("x:", x)
// x: tensor([-2., -1.,  0.,  1.,  2.], requires_grad=True)
// >>> print("y:", y)
// y: tensor([-1., -1.,  0.,  1.,  1.], grad_fn=<SignBackward0>)
// >>> # Compute the loss (just an example)
// >>> loss = y.sum()
// >>> # Backward pass: Compute the gradients
// >>> loss.backward()
// >>> print("\nBackward pass:")
// Backward pass:
// >>> print("x.grad:", x.grad)
// x.grad: tensor([0., 0., 0., 0., 0.])
    fn should_diff_sign() {
    type FT = FloatElem<TestBackend>;
    fn should_diff_sin() {
// Metal has less precise trigonometric functions
    type FT = FloatElem<TestBackend>;
    fn should_diff_matmul_with_slice() {
    fn should_diff_matmul_with_slice_assign() {
    fn should_diff_matmul_with_slice_assign_complex() {
    fn slice_assign_diff_should_give_same_results_as_cat() {
    type FT = FloatElem<TestBackend>;
    fn test_softmax_grad() {
    fn test_log_softmax_grad() {
// f16 gradients from log-softmax + matmul amplify error, so we we increase the tolerance
// to account for limited precision and large representable step sizes in this range.
    fn test_quiet_softmax_grad() {
// Precision is quite bad yet on softmax grad especially with half precision.
    type FT = FloatElem<TestBackend>;
    fn should_diff_sort() {
    fn should_diff_sort_with_indices() {
    type FT = FloatElem<TestBackend>;
    fn should_diff_sqrt() {
    fn should_diff_sub() {
    fn should_diff_sub_scalar() {
    fn test_sub_complex_1() {
    type FT = FloatElem<TestBackend>;
    fn should_diff_tanh() {
    type FT = FloatElem<TestBackend>;
    fn should_diff_transpose() {
    fn should_diff_swap_dims() {
/// Duplicate the given object for each node that requires gradients.
///
/// # Notes
///
/// This is useful since you don't have to keep N cloned references alive event if just 1 node
/// will be updated.
///
/// If the object is a tensor and if one reference exists, it can be updated inplace.
pub fn duplicate<T: Clone + core::fmt::Debug, const N: usize>(
    nodes: &[Option<NodeRef>;
/// Tensor backend that uses the [candle](candle_core) crate for executing tensor operations.
///
/// It is compatible with a wide range of hardware configurations, including CPUs and GPUs
/// that support CUDA or Metal. Additionally, the backend can be compiled to `wasm` when using the CPU.
pub struct Candle<F = f32, I = i64>
where
    F: FloatCandleElement,
    I: IntCandleElement,
{
/// The device type for the candle backend.
/// The device struct when using the `candle` backend.
///
/// To create a Cuda or Metal device from the index, use the associated methods to create the variant:
/// ```no_run
/// use burn_candle::CandleDevice;
///
/// // Create a Cuda device from its index
/// let device = CandleDevice::cuda(0);
/// // Create a Metal device from its index
/// let device = CandleDevice::metal(0);
/// ```
pub enum CandleDevice {
/// CPU device.
/// Cuda device with the given index. The index is the index of the Cuda device in the list of
/// all Cuda devices found on the system.
/// Metal device with the given index. The index is the index of the Metal device in the list of
/// all Metal devices found on the system.
impl CandleDevice {
/// Create a Cuda device with the given index.
/// The index is the index of the Cuda device in the list of all Cuda devices found on the system.
    pub fn cuda(index: usize) -> Self {
/// Create a Metal device with the given index.
/// The index is the index of the Metal device in the list of all Metal devices found on the system.
    pub fn metal(index: usize) -> Self {
/// A Cuda device for the `candle` backend.
pub struct CudaDevice {
/// The index of the Cuda device in the list of all devices on the system.
impl PartialEq for CudaDevice {
    fn eq(&self, other: &Self) -> bool {
impl Eq for CudaDevice {
/// A Metal device for the `candle` backend.
pub struct MetalDevice {
/// The index of the Metal device in the list of all devices on the system.
impl PartialEq for MetalDevice {
    fn eq(&self, other: &Self) -> bool {
impl Eq for MetalDevice {
impl From<CandleDevice> for candle_core::Device {
    fn from(device: CandleDevice) -> Self {
impl From<candle_core::Device> for CandleDevice {
    fn from(device: candle_core::Device) -> Self {
impl DeviceOps for CandleDevice {
    fn id(&self) -> burn_tensor::backend::DeviceId {
impl Default for CandleDevice {
    fn default() -> Self {
impl<F: FloatCandleElement, I: IntCandleElement> Backend for Candle<F, I> {
    type Device = CandleDevice;
    type FloatTensorPrimitive = CandleTensor;
    type FloatElem = F;
    type IntTensorPrimitive = CandleTensor;
    type IntElem = I;
    type BoolTensorPrimitive = CandleTensor;
    type BoolElem = u8;
    type QuantizedTensorPrimitive = CandleQTensor;
    type QuantizedEncoding = u8;
    fn ad_enabled() -> bool {
    fn name(device: &Self::Device) -> String {
    fn seed(seed: u64) {
// TODO submit an issue at Candle
    fn sync(device: &Device<Self>) {
// For some reason, device.wait_until_completed() does not seem to work,
// and neither does writing and reading a value with into_data
/// Candle element
pub trait CandleElement: Element + WithDType {
/// Candle float element
pub trait FloatCandleElement: CandleElement + FloatDType {
/// Candle int element
pub trait IntCandleElement: CandleElement {
impl CandleElement for f64 {
impl FloatCandleElement for f64 {
impl CandleElement for f32 {
impl FloatCandleElement for f32 {
impl CandleElement for f16 {
impl FloatCandleElement for f16 {
impl CandleElement for bf16 {
impl FloatCandleElement for bf16 {
impl CandleElement for u8 {
impl IntCandleElement for u8 {
//maybe also a quantization type?
impl CandleElement for u32 {
impl IntCandleElement for u32 {
impl CandleElement for i64 {
impl IntCandleElement for i64 {
// TODO remove when backend filled
//! Burn Candle Backend
    pub type TestBackend = Candle<f32, i64>;
    pub type ReferenceBackend = burn_tch::LibTorch<f32>;
    pub type TestTensor<const D: usize> = burn_tensor::Tensor<TestBackend, D>;
    pub type ReferenceTensor<const D: usize> = burn_tensor::Tensor<ReferenceBackend, D>;
    pub type TestTensorInt<const D: usize> = burn_tensor::Tensor<TestBackend, D, burn_tensor::Int>;
    pub type TestTensorBool<const D: usize> =
        burn_tensor::Tensor<TestBackend, D, burn_tensor::Bool>;
    type TestAutodiffBackend = burn_autodiff::Autodiff<TestBackend>;
    type TestAutodiffTensor<const D: usize> = burn_tensor::Tensor<TestAutodiffBackend, D>;
    pub type FloatType = f32;
// test activation
// test module
// burn_tensor::testgen_module_conv2d!();
// burn_tensor::testgen_module_conv_transpose1d!();
// burn_tensor::testgen_module_conv_transpose2d!();
// burn_tensor::testgen_module_max_pool1d!();
// burn_tensor::testgen_module_max_pool2d!();
// burn_tensor::testgen_module_avg_pool1d!();
// burn_tensor::testgen_module_avg_pool2d!();
// burn_tensor::testgen_module_adaptive_avg_pool1d!();
// burn_tensor::testgen_module_adaptive_avg_pool2d!();
// test ops
// burn_tensor::testgen_aggregation!();
// burn_tensor::testgen_div!();
// commented out due to macos CI failure, see #2427
// burn_tensor::testgen_remainder!();
// TODO: https://github.com/tracel-ai/burn/issues/1237
//
// burn_tensor::testgen_powf_scalar!();
// burn_tensor::testgen_powf!();
// test stats
// Behavior
// burn_autodiff::testgen_ad_broadcast!();
// Activation
// Modules
// burn_autodiff::testgen_ad_conv1d!();
// burn_autodiff::testgen_ad_conv2d!();
// burn_autodiff::testgen_ad_conv_transpose1d!();
// burn_autodiff::testgen_ad_conv_transpose2d!();
// burn_autodiff::testgen_ad_max_pool1d!();
// burn_autodiff::testgen_ad_max_pool2d!();
// burn_autodiff::testgen_ad_avg_pool1d!();
// burn_autodiff::testgen_ad_avg_pool2d!();
// burn_autodiff::testgen_ad_adaptive_avg_pool1d!();
// burn_autodiff::testgen_ad_adaptive_avg_pool2d!();
// Tensor
// burn_autodiff::testgen_ad_cat!();
// commented out due to macos CI failure, see #2427
// burn_autodiff::testgen_ad_remainder!();
impl<F: FloatCandleElement, I: IntCandleElement> ActivationOps<Self> for Candle<F, I> {
    fn gelu(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
    fn relu(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
pub fn cat(tensors: Vec<CandleTensor>, dim: usize) -> CandleTensor {
pub fn from_data<E: CandleElement>(data: TensorData, device: &CandleDevice) -> CandleTensor {
pub fn into_data(tensor: CandleTensor) -> TensorData {
    fn tensor_data_from_dtype<T: WithDType + Element>(tensor: &CandleTensor) -> TensorData {
pub fn to_device(tensor: CandleTensor, device: &CandleDevice) -> CandleTensor {
pub fn empty(shape: Shape, device: &CandleDevice, dtype: candle_core::DType) -> CandleTensor {
pub fn swap_dims(mut tensor: CandleTensor, dim1: usize, dim2: usize) -> CandleTensor {
pub fn permute(tensor: CandleTensor, axes: &[usize]) -> CandleTensor {
pub fn flip(tensor: CandleTensor, axes: &[usize]) -> CandleTensor {
// FIXME: Replace with an appropriate method when Candle provides one.
pub fn reshape(tensor: CandleTensor, shape: Shape) -> CandleTensor {
pub fn device(tensor: &CandleTensor) -> CandleDevice {
pub fn shape(tensor: &CandleTensor) -> Shape {
pub fn slice(tensor: CandleTensor, ranges: &[std::ops::Range<usize>]) -> CandleTensor {
pub fn slice_assign(
    tensor: CandleTensor,
    ranges: &[std::ops::Range<usize>],
    value: CandleTensor,
) -> CandleTensor {
pub fn narrow(tensor: CandleTensor, dim: usize, start: usize, length: usize) -> CandleTensor {
pub fn chunk(tensor: CandleTensor, chunks: usize, dim: usize) -> Vec<CandleTensor> {
pub fn expand(tensor: CandleTensor, shape: Shape) -> CandleTensor {
pub fn sign(tensor: CandleTensor) -> CandleTensor {
pub fn mask_where_broadcasted(
    tensor: CandleTensor,
    mask: CandleTensor,
    value: CandleTensor,
) -> CandleTensor {
impl<F: FloatCandleElement, I: IntCandleElement> BoolTensorOps<Self> for Candle<F, I> {
    fn bool_empty(shape: Shape, device: &Device<Self>) -> BoolTensor<Self> {
    async fn bool_into_data(tensor: BoolTensor<Self>) -> TensorData {
    fn bool_from_data(data: TensorData, device: &Device<Self>) -> BoolTensor<Self> {
    fn bool_into_int(tensor: BoolTensor<Self>) -> IntTensor<Self> {
    fn bool_into_float(tensor: BoolTensor<Self>) -> FloatTensor<Self> {
    fn bool_device(tensor: &BoolTensor<Self>) -> Device<Self> {
    fn bool_to_device(tensor: BoolTensor<Self>, device: &Device<Self>) -> BoolTensor<Self> {
    fn bool_reshape(tensor: BoolTensor<Self>, shape: Shape) -> BoolTensor<Self> {
    fn bool_slice(tensor: BoolTensor<Self>, ranges: &[std::ops::Range<usize>]) -> BoolTensor<Self> {
    fn bool_slice_assign(
        tensor: BoolTensor<Self>,
        ranges: &[std::ops::Range<usize>],
        value: BoolTensor<Self>,
    ) -> BoolTensor<Self> {
    fn bool_cat(tensors: Vec<BoolTensor<Self>>, dim: usize) -> BoolTensor<Self> {
    fn bool_equal(lhs: BoolTensor<Self>, rhs: BoolTensor<Self>) -> BoolTensor<Self> {
    fn bool_not(tensor: BoolTensor<Self>) -> BoolTensor<Self> {
    fn bool_and(lhs: BoolTensor<Self>, rhs: BoolTensor<Self>) -> BoolTensor<Self> {
    fn bool_or(lhs: BoolTensor<Self>, rhs: BoolTensor<Self>) -> BoolTensor<Self> {
    fn bool_swap_dims(tensor: BoolTensor<Self>, dim1: usize, dim2: usize) -> BoolTensor<Self> {
    fn bool_permute(tensor: BoolTensor<Self>, axes: &[usize]) -> BoolTensor<Self> {
    fn bool_flip(tensor: BoolTensor<Self>, axes: &[usize]) -> BoolTensor<Self> {
    fn bool_expand(tensor: BoolTensor<Self>, shape: Shape) -> BoolTensor<Self> {
pub(crate) fn fill<E: CandleElement, S: Into<Shape>>(
    value: E,
    shape: S,
    dtype: DType,
    device: &Device,
) -> Tensor {
pub(crate) fn fill_like<E: CandleElement>(value: E, reference_tensor: &Tensor) -> Tensor {
impl<F: FloatCandleElement, I: IntCandleElement> IntTensorOps<Self> for Candle<F, I> {
    fn int_empty(shape: Shape, device: &Device<Self>) -> IntTensor<Self> {
    async fn int_into_data(tensor: IntTensor<Self>) -> TensorData {
    fn int_from_data(data: TensorData, device: &Device<Self>) -> IntTensor<Self> {
    fn int_device(tensor: &IntTensor<Self>) -> Device<Self> {
    fn int_to_device(tensor: IntTensor<Self>, device: &Device<Self>) -> IntTensor<Self> {
    fn int_reshape(tensor: IntTensor<Self>, shape: Shape) -> IntTensor<Self> {
    fn int_slice(tensor: IntTensor<Self>, indices: &[std::ops::Range<usize>]) -> IntTensor<Self> {
    fn int_slice_assign(
        tensor: IntTensor<Self>,
        indices: &[std::ops::Range<usize>],
        value: IntTensor<Self>,
    ) -> IntTensor<Self> {
    fn int_into_float(tensor: IntTensor<Self>) -> FloatTensor<Self> {
    fn int_mask_where(
        tensor: IntTensor<Self>,
        mask: BoolTensor<Self>,
        source: IntTensor<Self>,
    ) -> IntTensor<Self> {
    fn int_mask_fill(
        tensor: IntTensor<Self>,
        mask: BoolTensor<Self>,
        value: IntElem<Self>,
    ) -> IntTensor<Self> {
    fn int_gather(
        dim: usize,
        tensor: IntTensor<Self>,
        indices: IntTensor<Self>,
    ) -> IntTensor<Self> {
    fn int_scatter(
        dim: usize,
        tensor: IntTensor<Self>,
        indices: IntTensor<Self>,
        value: IntTensor<Self>,
    ) -> IntTensor<Self> {
    fn int_select(
        tensor: IntTensor<Self>,
        dim: usize,
        indices: IntTensor<Self>,
    ) -> IntTensor<Self> {
    fn int_select_assign(
        tensor: IntTensor<Self>,
        dim: usize,
        indices: IntTensor<Self>,
        value: IntTensor<Self>,
    ) -> IntTensor<Self> {
    fn int_cat(tensors: Vec<IntTensor<Self>>, dim: usize) -> IntTensor<Self> {
    fn int_equal(lhs: IntTensor<Self>, rhs: IntTensor<Self>) -> BoolTensor<Self> {
    fn int_equal_elem(lhs: IntTensor<Self>, rhs: IntElem<Self>) -> BoolTensor<Self> {
    fn int_greater(lhs: IntTensor<Self>, rhs: IntTensor<Self>) -> BoolTensor<Self> {
    fn int_greater_elem(lhs: IntTensor<Self>, rhs: IntElem<Self>) -> BoolTensor<Self> {
    fn int_greater_equal(lhs: IntTensor<Self>, rhs: IntTensor<Self>) -> BoolTensor<Self> {
    fn int_greater_equal_elem(lhs: IntTensor<Self>, rhs: IntElem<Self>) -> BoolTensor<Self> {
    fn int_lower(lhs: IntTensor<Self>, rhs: IntTensor<Self>) -> BoolTensor<Self> {
    fn int_lower_elem(lhs: IntTensor<Self>, rhs: IntElem<Self>) -> BoolTensor<Self> {
    fn int_lower_equal(lhs: IntTensor<Self>, rhs: IntTensor<Self>) -> BoolTensor<Self> {
    fn int_lower_equal_elem(lhs: IntTensor<Self>, rhs: IntElem<Self>) -> BoolTensor<Self> {
    fn int_add(lhs: IntTensor<Self>, rhs: IntTensor<Self>) -> IntTensor<Self> {
    fn int_add_scalar(lhs: IntTensor<Self>, rhs: IntElem<Self>) -> IntTensor<Self> {
    fn int_sub(lhs: IntTensor<Self>, rhs: IntTensor<Self>) -> IntTensor<Self> {
    fn int_sub_scalar(lhs: IntTensor<Self>, rhs: IntElem<Self>) -> IntTensor<Self> {
    fn int_mul(lhs: IntTensor<Self>, rhs: IntTensor<Self>) -> IntTensor<Self> {
    fn int_mul_scalar(lhs: IntTensor<Self>, rhs: IntElem<Self>) -> IntTensor<Self> {
    fn int_div(lhs: IntTensor<Self>, rhs: IntTensor<Self>) -> IntTensor<Self> {
    fn int_div_scalar(lhs: IntTensor<Self>, rhs: IntElem<Self>) -> IntTensor<Self> {
// Candle implements scalar a/b as a * (1/b). With ints 1/b is rounded to 0 so we always obtain 0.
    fn int_remainder(lhs: IntTensor<Self>, rhs: IntTensor<Self>) -> IntTensor<Self> {
    fn int_remainder_scalar(lhs: IntTensor<Self>, rhs: IntElem<Self>) -> IntTensor<Self> {
// Same problem as int_div_scalar.
    fn int_zeros(shape: Shape, device: &Device<Self>) -> IntTensor<Self> {
    fn int_ones(shape: Shape, device: &Device<Self>) -> IntTensor<Self> {
    fn int_sum(tensor: IntTensor<Self>) -> IntTensor<Self> {
    fn int_sum_dim(tensor: IntTensor<Self>, dim: usize) -> IntTensor<Self> {
    fn int_prod(tensor: IntTensor<Self>) -> IntTensor<Self> {
    fn int_prod_dim(tensor: IntTensor<Self>, dim: usize) -> IntTensor<Self> {
    fn int_mean_dim(tensor: IntTensor<Self>, dim: usize) -> IntTensor<Self> {
// Candle implements scalar a/b as a * (1/b). With ints 1/b is rounded to 0 so we always obtain 0.
    fn int_argmax(tensor: IntTensor<Self>, dim: usize) -> IntTensor<Self> {
    fn int_argmin(tensor: IntTensor<Self>, dim: usize) -> IntTensor<Self> {
    fn int_abs(tensor: IntTensor<Self>) -> IntTensor<Self> {
// Ugly type conversion here as Candle does not support unary ops on ints
    fn int_swap_dims(tensor: IntTensor<Self>, dim1: usize, dim2: usize) -> IntTensor<Self> {
    fn int_random(
        shape: Shape,
        distribution: Distribution,
        device: &Device<Self>,
    ) -> IntTensor<Self> {
    fn int_permute(tensor: IntTensor<Self>, axes: &[usize]) -> IntTensor<Self> {
    fn int_flip(tensor: IntTensor<Self>, axes: &[usize]) -> IntTensor<Self> {
    fn int_expand(tensor: IntTensor<Self>, shape: Shape) -> IntTensor<Self> {
    fn int_sign(tensor: IntTensor<Self>) -> IntTensor<Self> {
    fn bitwise_and(lhs: IntTensor<Self>, rhs: IntTensor<Self>) -> IntTensor<Self> {
    fn bitwise_and_scalar(lhs: IntTensor<Self>, rhs: IntElem<Self>) -> IntTensor<Self> {
    fn bitwise_or(lhs: IntTensor<Self>, rhs: IntTensor<Self>) -> IntTensor<Self> {
    fn bitwise_or_scalar(lhs: IntTensor<Self>, rhs: IntElem<Self>) -> IntTensor<Self> {
    fn bitwise_xor(lhs: IntTensor<Self>, rhs: IntTensor<Self>) -> IntTensor<Self> {
    fn bitwise_xor_scalar(lhs: IntTensor<Self>, rhs: IntElem<Self>) -> IntTensor<Self> {
    fn bitwise_not(tensor: IntTensor<Self>) -> IntTensor<Self> {
    fn bitwise_left_shift(lhs: IntTensor<Self>, rhs: IntTensor<Self>) -> IntTensor<Self> {
    fn bitwise_right_shift(lhs: IntTensor<Self>, rhs: IntTensor<Self>) -> IntTensor<Self> {
    fn bitwise_left_shift_scalar(lhs: IntTensor<Self>, rhs: IntElem<Self>) -> IntTensor<Self> {
    fn bitwise_right_shift_scalar(lhs: IntTensor<Self>, rhs: IntElem<Self>) -> IntTensor<Self> {
impl<F: FloatCandleElement, I: IntCandleElement> ModuleOps<Self> for Candle<F, I> {
    fn conv1d(
        x: FloatTensor<Self>,
        weight: FloatTensor<Self>,
        bias: Option<FloatTensor<Self>>,
        options: ConvOptions<1>,
    ) -> FloatTensor<Self> {
    fn conv2d(
        x: FloatTensor<Self>,
        weight: FloatTensor<Self>,
        bias: Option<FloatTensor<Self>>,
        options: ConvOptions<2>,
    ) -> FloatTensor<Self> {
    fn deform_conv2d(
        x: FloatTensor<Self>,
        offset: FloatTensor<Self>,
        weight: FloatTensor<Self>,
        mask: Option<FloatTensor<Self>>,
        bias: Option<FloatTensor<Self>>,
        options: DeformConvOptions<2>,
    ) -> FloatTensor<Self> {
    fn deform_conv2d_backward(
        x: FloatTensor<Self>,
        offset: FloatTensor<Self>,
        weight: FloatTensor<Self>,
        mask: Option<FloatTensor<Self>>,
        bias: Option<FloatTensor<Self>>,
        output_grad: FloatTensor<Self>,
        options: DeformConvOptions<2>,
    ) -> DeformConv2dBackward<Self> {
    fn conv3d(
        x: FloatTensor<Self>,
        weight: FloatTensor<Self>,
        bias: Option<FloatTensor<Self>>,
        options: ConvOptions<3>,
    ) -> FloatTensor<Self> {
    fn conv_transpose1d(
        x: FloatTensor<Self>,
        weight: FloatTensor<Self>,
        bias: Option<FloatTensor<Self>>,
        options: ConvTransposeOptions<1>,
    ) -> FloatTensor<Self> {
    fn conv_transpose2d(
        x: FloatTensor<Self>,
        weight: FloatTensor<Self>,
        bias: Option<FloatTensor<Self>>,
        options: ConvTransposeOptions<2>,
    ) -> FloatTensor<Self> {
    fn conv_transpose3d(
        x: FloatTensor<Self>,
        weight: FloatTensor<Self>,
        bias: Option<FloatTensor<Self>>,
        options: ConvTransposeOptions<3>,
    ) -> FloatTensor<Self> {
    fn avg_pool2d(
        x: FloatTensor<Self>,
        kernel_size: [usize;
    fn avg_pool2d_backward(
        x: FloatTensor<Self>,
        grad: FloatTensor<Self>,
        kernel_size: [usize;
    fn max_pool2d(
        x: FloatTensor<Self>,
        kernel_size: [usize;
    fn max_pool2d_with_indices(
        x: FloatTensor<Self>,
        kernel_size: [usize;
    fn max_pool2d_with_indices_backward(
        x: FloatTensor<Self>,
        kernel_size: [usize;
    fn adaptive_avg_pool2d(x: FloatTensor<Self>, output_size: [usize;
    fn adaptive_avg_pool2d_backward(
        x: FloatTensor<Self>,
        grad: FloatTensor<Self>,
    ) -> FloatTensor<Self> {
    fn interpolate(
        x: FloatTensor<Self>,
        output_size: [usize;
    fn interpolate_backward(
        x: FloatTensor<Self>,
        grad: FloatTensor<Self>,
        output_size: [usize;
impl<F: FloatCandleElement, I: IntCandleElement> QTensorOps<Self> for Candle<F, I> {
    fn q_from_data(data: TensorData, device: &Device<Self>) -> QuantizedTensor<Self> {
// no i8 support
    fn quantize(
        _tensor: FloatTensor<Self>,
        _scheme: &QuantScheme,
        _qparams: QuantizationParametersPrimitive<Self>,
    ) -> QuantizedTensor<Self> {
    fn dequantize(_tensor: QuantizedTensor<Self>) -> FloatTensor<Self> {
    fn q_device(tensor: &QuantizedTensor<Self>) -> Device<Self> {
    fn q_to_device(
        _tensor: QuantizedTensor<Self>,
        _device: &Device<Self>,
    ) -> QuantizedTensor<Self> {
    fn q_reshape(tensor: QuantizedTensor<Self>, shape: Shape) -> QuantizedTensor<Self> {
    async fn q_into_data(tensor: QuantizedTensor<Self>) -> TensorData {
    fn q_swap_dims(
        _tensor: QuantizedTensor<Self>,
        _dim1: usize,
        _dim2: usize,
    ) -> QuantizedTensor<Self> {
    fn q_permute(_tensor: QuantizedTensor<Self>, _axes: &[usize]) -> QuantizedTensor<Self> {
    fn q_flip(_tensor: QuantizedTensor<Self>, _axes: &[usize]) -> QuantizedTensor<Self> {
    fn q_gather(
        _dim: usize,
        _tensor: QuantizedTensor<Self>,
        _indices: IntTensor<Self>,
    ) -> QuantizedTensor<Self> {
    fn q_select(
        _tensor: QuantizedTensor<Self>,
        _dim: usize,
        _indices: IntTensor<Self>,
    ) -> QuantizedTensor<Self> {
    fn q_slice(_tensor: QuantizedTensor<Self>, _ranges: &[Range<usize>]) -> QuantizedTensor<Self> {
    fn q_expand(_tensor: QuantizedTensor<Self>, _shape: Shape) -> QuantizedTensor<Self> {
impl<F: FloatCandleElement, I: IntCandleElement> FloatTensorOps<Self> for Candle<F, I> {
    fn float_from_data(data: TensorData, device: &Device<Self>) -> CandleTensor {
    fn float_random(
        shape: Shape,
        distribution: Distribution,
        device: &Device<Self>,
    ) -> FloatTensor<Self> {
    async fn float_into_data(tensor: CandleTensor) -> TensorData {
    fn float_device(tensor: &CandleTensor) -> Device<Self> {
    fn float_to_device(tensor: CandleTensor, device: &Device<Self>) -> CandleTensor {
    fn float_into_int(tensor: CandleTensor) -> IntTensor<Self> {
    fn float_empty(shape: Shape, device: &Device<Self>) -> FloatTensor<Self> {
    fn float_add(lhs: FloatTensor<Self>, rhs: FloatTensor<Self>) -> FloatTensor<Self> {
    fn float_add_scalar(lhs: FloatTensor<Self>, rhs: FloatElem<Self>) -> FloatTensor<Self> {
    fn float_sub(lhs: FloatTensor<Self>, rhs: FloatTensor<Self>) -> FloatTensor<Self> {
    fn float_sub_scalar(lhs: FloatTensor<Self>, rhs: FloatElem<Self>) -> FloatTensor<Self> {
    fn float_mul(lhs: FloatTensor<Self>, rhs: FloatTensor<Self>) -> FloatTensor<Self> {
    fn float_mul_scalar(lhs: FloatTensor<Self>, rhs: FloatElem<Self>) -> FloatTensor<Self> {
    fn float_div(lhs: FloatTensor<Self>, rhs: FloatTensor<Self>) -> FloatTensor<Self> {
    fn float_div_scalar(lhs: FloatTensor<Self>, rhs: FloatElem<Self>) -> FloatTensor<Self> {
    fn float_remainder(lhs: FloatTensor<Self>, rhs: FloatTensor<Self>) -> FloatTensor<Self> {
    fn float_remainder_scalar(lhs: FloatTensor<Self>, rhs: FloatElem<Self>) -> FloatTensor<Self> {
// In PyTorch, remainder can also be defined as torch.remainder(a, b) == a - a.div(b, rounding_mode="floor") * b
    fn float_matmul(lhs: FloatTensor<Self>, rhs: FloatTensor<Self>) -> FloatTensor<Self> {
    fn float_swap_dims(tensor: FloatTensor<Self>, dim1: usize, dim2: usize) -> FloatTensor<Self> {
    fn float_reshape(tensor: FloatTensor<Self>, shape: Shape) -> FloatTensor<Self> {
    fn float_gather(
        dim: usize,
        tensor: FloatTensor<Self>,
        indices: IntTensor<Self>,
    ) -> FloatTensor<Self> {
    fn float_scatter(
        dim: usize,
        tensor: FloatTensor<Self>,
        indices: IntTensor<Self>,
        value: FloatTensor<Self>,
    ) -> FloatTensor<Self> {
    fn float_select(
        tensor: FloatTensor<Self>,
        dim: usize,
        indices: IntTensor<Self>,
    ) -> FloatTensor<Self> {
    fn float_select_assign(
        tensor: FloatTensor<Self>,
        dim: usize,
        indices: IntTensor<Self>,
        value: FloatTensor<Self>,
    ) -> FloatTensor<Self> {
    fn float_slice(
        tensor: FloatTensor<Self>,
        ranges: &[std::ops::Range<usize>],
    ) -> FloatTensor<Self> {
    fn float_slice_assign(
        tensor: FloatTensor<Self>,
        ranges: &[std::ops::Range<usize>],
        value: FloatTensor<Self>,
    ) -> FloatTensor<Self> {
    fn float_mask_where(
        tensor: FloatTensor<Self>,
        mask: BoolTensor<Self>,
        value: FloatTensor<Self>,
    ) -> FloatTensor<Self> {
    fn float_mask_fill(
        tensor: FloatTensor<Self>,
        mask: BoolTensor<Self>,
        value: FloatElem<Self>,
    ) -> FloatTensor<Self> {
    fn float_equal(lhs: FloatTensor<Self>, rhs: FloatTensor<Self>) -> BoolTensor<Self> {
    fn float_equal_elem(lhs: FloatTensor<Self>, rhs: FloatElem<Self>) -> BoolTensor<Self> {
    fn float_greater(lhs: FloatTensor<Self>, rhs: FloatTensor<Self>) -> BoolTensor<Self> {
    fn float_greater_elem(lhs: FloatTensor<Self>, rhs: FloatElem<Self>) -> BoolTensor<Self> {
    fn float_greater_equal(lhs: FloatTensor<Self>, rhs: FloatTensor<Self>) -> BoolTensor<Self> {
    fn float_greater_equal_elem(lhs: FloatTensor<Self>, rhs: FloatElem<Self>) -> BoolTensor<Self> {
    fn float_lower(lhs: FloatTensor<Self>, rhs: FloatTensor<Self>) -> BoolTensor<Self> {
    fn float_lower_elem(lhs: FloatTensor<Self>, rhs: FloatElem<Self>) -> BoolTensor<Self> {
    fn float_lower_equal(lhs: FloatTensor<Self>, rhs: FloatTensor<Self>) -> BoolTensor<Self> {
    fn float_lower_equal_elem(lhs: FloatTensor<Self>, rhs: FloatElem<Self>) -> BoolTensor<Self> {
    fn float_sum(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
    fn float_sum_dim(tensor: FloatTensor<Self>, dim: usize) -> FloatTensor<Self> {
    fn float_mean_dim(tensor: FloatTensor<Self>, dim: usize) -> FloatTensor<Self> {
    fn float_exp(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
    fn float_log(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
    fn float_log1p(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
    fn float_powf_scalar(tensor: FloatTensor<Self>, value: f32) -> FloatTensor<Self> {
    fn float_sqrt(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
    fn float_abs(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
    fn float_cos(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
    fn float_sin(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
    fn float_tanh(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
    fn float_round(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
// implements round_to_even for consistent behavior vs libtorch
// https://github.com/pytorch/pytorch/blob/main/torch/csrc/jit/runtime/register_ops_utils.h#L65-L67
    fn float_floor(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
    fn float_ceil(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
    fn float_erf(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
    fn float_cat(tensors: Vec<FloatTensor<Self>>, dim: usize) -> FloatTensor<Self> {
    fn float_argmax(tensor: FloatTensor<Self>, dim: usize) -> IntTensor<Self> {
    fn float_argmin(tensor: FloatTensor<Self>, dim: usize) -> IntTensor<Self> {
    fn float_clamp_max(tensor: FloatTensor<Self>, max: FloatElem<Self>) -> FloatTensor<Self> {
    fn float_clamp_min(tensor: FloatTensor<Self>, min: FloatElem<Self>) -> FloatTensor<Self> {
    fn float_clamp(
        tensor: FloatTensor<Self>,
        min: FloatElem<Self>,
        max: FloatElem<Self>,
    ) -> FloatTensor<Self> {
    fn float_recip(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
    fn float_powf(lhs: FloatTensor<Self>, rhs: FloatTensor<Self>) -> FloatTensor<Self> {
//broadcast_pow is in main but not yet published
//note: probably replace once pow once 0.3.3 is out
//see: https://github.com/huggingface/candle/pull/1583/files#diff-6319fa1e16dadc4c7b4e25698139703d93b70f30a1f8e2ac0999978e39efaa81R2594
    fn float_permute(tensor: FloatTensor<Self>, axes: &[usize]) -> FloatTensor<Self> {
    fn float_flip(tensor: FloatTensor<Self>, axes: &[usize]) -> FloatTensor<Self> {
    fn float_expand(tensor: FloatTensor<Self>, shape: Shape) -> FloatTensor<Self> {
    fn float_sign(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
    fn float_cast(tensor: FloatTensor<Self>, dtype: FloatDType) -> FloatTensor<Self> {
impl<F: FloatCandleElement, I: IntCandleElement> TransactionOps<Self> for Candle<F, I> {
/// A tensor that uses the candle backend.
pub struct CandleTensor {
impl TensorMetadata for CandleTensor {
    fn dtype(&self) -> DType {
    fn shape(&self) -> Shape {
impl CandleTensor {
/// Create a new tensor.
    pub fn new(tensor: candle_core::Tensor) -> Self {
/// Creates a new tensor from data and a device.
///
/// # Arguments
///
/// * `data` - The tensor's data.
/// * `device` - The device on which the tensor will be allocated.
///
/// # Returns
///
/// A new tensor.
    pub fn from_data<E: CandleElement>(data: TensorData, device: CandleDevice) -> Self {
/// A quantized tensor for the candle backend.
pub struct CandleQTensor {
/// The quantized tensor.
// NOTE: candle  does not implement `WithDType` for i8
/// The quantization scheme.
impl QTensorPrimitive for CandleQTensor {
    fn scheme(&self) -> &QuantScheme {
impl TensorMetadata for CandleQTensor {
    fn dtype(&self) -> DType {
    fn shape(&self) -> Shape {
/// Configuration IO error.
pub enum ConfigError {
/// Invalid format.
/// File not found.
impl core::fmt::Display for ConfigError {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
impl core::error::Error for ConfigError {
/// Configuration trait.
pub trait Config: Debug + serde::Serialize + serde::de::DeserializeOwned {
/// Saves the configuration to a file.
///
/// # Arguments
///
/// * `file` - File to save the configuration to.
///
/// # Returns
///
/// The output of the save operation.
    fn save<P: AsRef<std::path::Path>>(&self, file: P) -> std::io::Result<()> {
/// Loads the configuration from a file.
///
/// # Arguments
///
/// * `file` - File to load the configuration from.
///
/// # Returns
///
/// The loaded configuration.
    fn load<P: AsRef<std::path::Path>>(file: P) -> Result<Self, ConfigError> {
/// Loads the configuration from a binary buffer.
///
/// # Arguments
///
/// * `data` - Binary buffer to load the configuration from.
///
/// # Returns
///
/// The loaded configuration.
    fn load_binary(data: &[u8]) -> Result<Self, ConfigError> {
/// Converts a configuration to a JSON string.
///
/// # Arguments
///
/// * `config` - Configuration to convert.
///
/// # Returns
///
/// The JSON string.
pub fn config_to_json<C: Config>(config: &C) -> String {
fn config_from_str<C: Config>(content: &str) -> Result<C, ConfigError> {
/// A progress struct that can be used to track the progress of a data loader.
pub struct Progress {
/// The number of items that have been processed.
/// The total number of items that need to be processed.
/// A data loader iterator that can be used to iterate over a data loader.
pub trait DataLoaderIterator<O>: Iterator<Item = O> {
/// Returns the progress of the data loader.
    fn progress(&self) -> Progress;
/// A data loader that can be used to iterate over a dataset.
pub trait DataLoader<B: Backend, O>: Send + Sync {
/// Returns a boxed [iterator](DataLoaderIterator) to iterate over the data loader.
    fn iter<'a>(&'a self) -> Box<dyn DataLoaderIterator<O> + 'a>;
/// The number of items (not the number of batches nor the number of iterations),
/// corresponding to the items_total of the progress returned by the iterator.
    fn num_items(&self) -> usize;
/// Move the data loader to the given device, ensuring the batches are assigned to the correct device.
    fn to_device(&self, device: &B::Device) -> Arc<dyn DataLoader<B, O>>;
/// Returns a new data loader containing a subset of the data.
///
/// The subset includes items from `start` (inclusive) to `end` (exclusive),
/// preserving the batch size and ordering of the original data loader.
///
/// # Arguments
///
/// * `start` - The starting index of the subset (inclusive).
/// * `end` - The ending index of the subset (exclusive).
///
/// # Returns
///
/// A boxed [`DataLoader`] instance containing only the specified range.
    fn slice(&self, start: usize, end: usize) -> Arc<dyn DataLoader<B, O>>;
/// A data loader that can be used to iterate over a dataset in batches.
pub struct BatchDataLoader<B: Backend, I, O> {
impl<B: Backend, I, O> Clone for BatchDataLoader<B, I, O> {
    fn clone(&self) -> Self {
impl<B: Backend, I, O> BatchDataLoader<B, I, O> {
/// Creates a new batch data loader.
///
/// # Arguments
///
/// * `strategy` - The batch strategy.
/// * `dataset` - The dataset.
/// * `batcher` - The batcher.
/// * `device`  - The device to use when loading a batch.
/// * `rng`     - The rng determining if the dataset is shuffled each time a dataloader
///   iterator is created.
///
/// # Returns
///
/// The batch data loader.
    pub fn new(
        strategy: Box<dyn BatchStrategy<I>>,
        dataset: Arc<dyn Dataset<I>>,
        batcher: Arc<dyn Batcher<B, I, O>>,
        device: B::Device,
        rng: Option<rand::rngs::StdRng>,
    ) -> Self {
/// A data loader iterator that can be used to iterate over a data loader.
struct BatchDataloaderIterator<B: Backend, I, O> {
impl<B, I, O> DataLoader<B, O> for BatchDataLoader<B, I, O>
where
    B: Backend,
    I: Send + Sync + Clone + 'static,
    O: Send + 'static,
{
    fn iter<'a>(&'a self) -> Box<dyn DataLoaderIterator<O> + 'a> {
// When starting a new iteration, we first check if the dataloader was created with an rng,
// implying that we should shuffle the dataset beforehand, while advancing the current
// rng to ensure that each new iteration shuffles the dataset differently.
    fn num_items(&self) -> usize {
    fn to_device(&self, device: &B::Device) -> Arc<dyn DataLoader<B, O>> {
    fn slice(&self, start: usize, end: usize) -> Arc<dyn DataLoader<B, O>> {
impl<B: Backend, I, O> BatchDataloaderIterator<B, I, O> {
/// Creates a new batch data loader iterator.
///
/// # Arguments
///
/// * `strategy` - The batch strategy.
/// * `dataset` - The dataset.
/// * `batcher` - The batcher.
/// * `device`  - The device to use when loading a batch.
///
/// # Returns
///
/// The batch data loader iterator.
    pub fn new(
        strategy: Box<dyn BatchStrategy<I>>,
        dataset: Arc<dyn Dataset<I>>,
        batcher: Arc<dyn Batcher<B, I, O>>,
        device: B::Device,
    ) -> Self {
impl<B: Backend, I, O> Iterator for BatchDataloaderIterator<B, I, O> {
    type Item = O;
    fn next(&mut self) -> Option<O> {
impl<B: Backend, I, O> DataLoaderIterator<O> for BatchDataloaderIterator<B, I, O> {
    fn progress(&self) -> Progress {
    fn test_batch_dataloader() {
    fn test_batch_dataloader_slice() {
/// A trait for batching items of type `I` into items of type `O`.
pub trait Batcher<B: Backend, I, O>: Send + Sync {
/// Batches the given items on the specified device.
///
/// # Arguments
///
/// * `items` - The items to batch.
/// * `device` - The backend device to use.
///
/// # Returns
///
/// The batched items.
    fn batch(&self, items: Vec<I>, device: &B::Device) -> O;
/// Test batcher
pub struct TestBatcher;
impl<I> Batcher<TestBackend, I, Vec<I>> for TestBatcher {
    fn batch(&self, items: Vec<I>, _device: &<TestBackend as Backend>::Device) -> Vec<I> {
/// A builder for data loaders.
pub struct DataLoaderBuilder<B: Backend, I, O> {
impl<B, I, O> DataLoaderBuilder<B, I, O>
where
    B: Backend,
    I: Send + Sync + Clone + std::fmt::Debug + 'static,
    O: Send + Clone + std::fmt::Debug + 'static,
{
/// Creates a new data loader builder.
///
/// # Arguments
///
/// * `batcher` - The batcher.
///
/// # Returns
///
/// The data loader builder.
    pub fn new<Bt>(batcher: Bt) -> Self
    where
        Bt: Batcher<B, I, O> + 'static,
    {
/// Sets the batch size to a fix number.
///
/// The [fix batch strategy](FixBatchStrategy) will be used.
///
/// # Arguments
///
/// * `batch_size` - The batch size.
///
/// # Returns
///
/// The data loader builder.
    pub fn batch_size(mut self, batch_size: usize) -> Self {
/// Sets the seed for shuffling.
///
/// Each time the dataloader starts a new iteration, the dataset will be shuffled.
///
/// # Arguments
///
/// * `seed` - The seed.
///
/// # Returns
///
/// The data loader builder.
    pub fn shuffle(mut self, seed: u64) -> Self {
/// Sets the number of workers.
///
/// - `Some(0)` or `None`: the dataloader will run without work threads.
/// - `Some(n); n > 0`: the dataloader will run with `n` background threads.
///
/// A 1-worker threaded dataloader will run loads in a background thread,
/// while a 0-worker threaded dataloader will run loads in the main thread.
///
/// # Arguments
///
/// * `num_workers` - The number of workers.
///
/// # Returns
///
/// The data loader builder.
    pub fn num_workers(mut self, num_workers: usize) -> Self {
/// Sets the data loader device.
///
/// # Arguments
///
/// * `device` - The device to use when loading a batch.
///
/// # Returns
///
/// The data loader builder.
    pub fn set_device(mut self, device: B::Device) -> Self {
/// Builds the data loader.
///
/// # Arguments
///
/// * `dataset` - The dataset.
///
/// # Returns
///
/// The data loader.
    pub fn build<D>(self, dataset: D) -> Arc<dyn DataLoader<B, O>>
    where
        D: Dataset<I> + 'static,
    {
    struct TestBatcherDevice;
    impl<I> Batcher<TestBackend, I, TestDevice> for TestBatcherDevice {
        fn batch(&self, _items: Vec<I>, device: &TestDevice) -> TestDevice {
    type TestDevice = <TestBackend as Backend>::Device;
    fn test_dataloader_no_workers() {
        type TestDevice = <TestBackend as Backend>::Device;
    fn test_dataloader_default_device() {
    fn test_dataloader_slice_multi_device() {
// Only one device exists...
// For uneven split, the last dataloader (partial dataset) will have the remaining item
/// Module for batching items.
/// Module to split a dataloader.
const MAX_QUEUED_ITEMS: usize = 100;
/// A multi-threaded data loader that can be used to iterate over a dataset.
pub struct MultiThreadDataLoader<B: Backend, I, O> {
// Configuration parameters needed for initialization
// The lazily initialized data loaders
/// A message that can be sent between threads.
pub enum Message<O> {
/// A batch of items.
/// The thread is done.
struct MultiThreadsDataloaderIterator<O> {
impl<B: Backend, I, O> MultiThreadDataLoader<B, I, O>
where
    I: Send + Sync + Clone + 'static,
    O: Send + 'static,
{
/// Creates a new multi-threaded batch data loader.
///
/// # Arguments
///
/// * `strategy` - The batch strategy.
/// * `dataset` - The dataset.
/// * `batcher` - The batcher.
/// * `num_threads` - The number of threads.
/// * `device`  - The device to use when loading a batch.
/// * `rng`     - The rng determining if the dataset is shuffled each time a dataloader
///   iterator is created.
///
/// # Returns
///
/// The multi-threaded batch data loader.
    pub fn new(
        strategy: Box<dyn BatchStrategy<I>>,
        dataset: Arc<dyn Dataset<I>>,
        batcher: Arc<dyn Batcher<B, I, O>>,
        num_threads: usize,
        device: B::Device,
        rng: Option<rand::rngs::StdRng>,
    ) -> Self {
/// Force initialization if needed.
    fn initialize(&self) -> &[BatchDataLoader<B, I, O>] {
// Pre-shuffle the dataset before split if shuffle is enabled.
// This ensures that each thread gets a uniform random sample of the dataset.
// Create more rngs from the first one, one for each new dataloader.
impl<B: Backend, I, O> DataLoader<B, O> for MultiThreadDataLoader<B, I, O>
where
    I: Send + Sync + Clone + 'static,
    O: Send + 'static + std::fmt::Debug,
{
    fn iter<'a>(&'a self) -> Box<dyn DataLoaderIterator<O> + 'a> {
// This will initialize the loader if it hasn't been initialized yet
// The receiver is probably gone, no need to panic, just need to stop
// iterating.
// Same thing.
    fn num_items(&self) -> usize {
// For num_items, we can directly use the dataset size without
// necessarily initializing the full loader
    fn to_device(&self, device: &B::Device) -> Arc<dyn DataLoader<B, O>> {
    fn slice(&self, start: usize, end: usize) -> Arc<dyn DataLoader<B, O>> {
impl<O> MultiThreadsDataloaderIterator<O> {
    pub fn new(
        receiver: mpsc::Receiver<Message<O>>,
        workers: Vec<thread::JoinHandle<()>>,
        progresses: Vec<Progress>,
    ) -> Self {
impl<O: std::fmt::Debug> DataLoaderIterator<O> for MultiThreadsDataloaderIterator<O> {
    fn progress(&self) -> Progress {
impl<O: std::fmt::Debug> Iterator for MultiThreadsDataloaderIterator<O> {
    type Item = O;
    fn next(&mut self) -> Option<O> {
    fn test_multi_thread_batch_dataloader() {
    fn test_multi_thread_batch_dataloader_shuffle() {
// Items is a deliberately ordered dataset.
// Unshuffled multithreaded loader
// No rng means no shuffling.
// Since the dataset is not shuffled, we expect each batch to contain the same item.
// Shuffled multithreaded loader
// The rng enables shuffling.
// Since the dataset is shuffled, we expect to see all items.
    fn test_multi_thread_batch_dataloader_incomplete_batches() {
/// Splits a dataloader into multiple partial dataloaders (one per device).
pub fn split_dataloader<B: Backend, O>(
    dataloader: Arc<dyn DataLoader<B, O>>,
    devices: &[B::Device],
) -> Vec<Arc<dyn DataLoader<B, O>>> {
    fn test_split_batch_dataloader() {
        type TestDevice = <TestBackend as Backend>::Device;
        pub struct TestBatcher;
        impl<I> Batcher<TestBackend, I, (Vec<I>, TestDevice)> for TestBatcher {
            fn batch(&self, items: Vec<I>, device: &TestDevice) -> (Vec<I>, TestDevice) {
// Only one device exists...
/// A strategy to batch items.
pub trait BatchStrategy<I>: Send + Sync {
/// Adds an item to the strategy.
///
/// # Arguments
///
/// * `item` - The item to add.
    fn add(&mut self, item: I);
/// Batches the items.
///
/// # Arguments
///
/// * `force` - Whether to force batching.
///
/// # Returns
///
/// The batched items.
    fn batch(&mut self, force: bool) -> Option<Vec<I>>;
/// Creates a new strategy of the same type.
///
/// # Returns
///
/// The new strategy.
    fn clone_dyn(&self) -> Box<dyn BatchStrategy<I>>;
/// Returns the expected batch size for this strategy.
///
/// # Returns
///
/// The batch size, or None if the strategy doesn't have a fixed batch size.
    fn batch_size(&self) -> Option<usize>;
/// A strategy to batch items with a fixed batch size.
pub struct FixBatchStrategy<I> {
impl<I> FixBatchStrategy<I> {
/// Creates a new strategy to batch items with a fixed batch size.
///
/// # Arguments
///
/// * `batch_size` - The batch size.
///
/// # Returns
///
/// The strategy.
    pub fn new(batch_size: usize) -> Self {
impl<I: Send + Sync + 'static> BatchStrategy<I> for FixBatchStrategy<I> {
    fn add(&mut self, item: I) {
    fn batch(&mut self, force: bool) -> Option<Vec<I>> {
    fn clone_dyn(&self) -> Box<dyn BatchStrategy<I>> {
    fn batch_size(&self) -> Option<usize> {
/// Dataloader module.
/// Dataset module.
/// Network module.
//! The core crate of Burn.
/// Re-export serde for proc macros.
/// The configuration module.
/// Data module.
/// Module for the neural network module.
/// Module for the recorder.
/// Module for the tensor.
// Tensor at root: `burn::Tensor`
/// Module for visual operations
/// Backend for test cases
pub type TestBackend = burn_ndarray::NdArray<f32>;
/// Backend for test cases
pub type TestBackend = burn_tch::LibTorch<f32>;
/// Backend for test cases
pub type TestBackend = burn_wgpu::Wgpu;
/// Backend for test cases
pub type TestBackend = burn_cuda::Cuda;
/// Backend for test cases
pub type TestBackend = burn_rocm::Rocm;
/// Backend for autodiff test cases
pub type TestAutodiffBackend = burn_autodiff::Autodiff<TestBackend>;
/// Simple linear module.
    pub struct SimpleLinear<B: Backend> {
    impl<B: Backend> SimpleLinear<B> {
        pub fn new(in_features: usize, out_features: usize, device: &B::Device) -> Self {
//! Structs and macros used by most projects. Add `use
//! burn::prelude::*` to your code to quickly get started with
//! Burn.
/// Type alias to `Vec<B::Device>` which supports `no_std` environments, but automatically using
/// the `alloc` crate.
pub type Devices<B> = Vec<Device<B>>;
// At the moment, our plan is to continue experimenting with the macro internally and monitor its development.
// We may consider making it public in the future.
        struct Mapper;
        impl<B: Backend> ModuleMapper<B> for Mapper {
            fn map_float<const D: usize>(
                &mut self,
                param: Param<Tensor<B, D>>,
            ) -> Param<Tensor<B, D>> {
        struct Visitor<'a, B: Backend> {
        impl<'a, B: Backend> ModuleVisitor<B> for Visitor<'a, B> {
            fn visit_float<const D: usize>(&mut self, param: &Param<Tensor<B, D>>) {
/// Trait for all neural network modules.
///
/// Modules should be created using the [derive](burn_derive::Module) attribute.
/// This will make your module trainable, savable and loadable via
/// `state` and `load`.
///
/// # Example
///
/// A module should have a [backend](crate::tensor::backend::Backend) defined as a generic
/// parameter B. This will be used by the [derive](burn_derive::Module) attribute to generate the code
/// necessary to optimize and train the module on any backend.
///
/// ```rust, ignore
/// // Not necessary when using the burn crate directly.
/// use burn_core as burn;
///
/// use burn::{
///     module::Module,
///     nn::Linear,
///     tensor::Tensor,
///     tensor::backend::Backend,
/// };
///
/// #[derive(Module, Debug)]
/// struct MyModule<B: Backend> {
///   my_param: Linear<B>,
///   my_other_field: usize,
/// }
/// ```
pub trait Module<B: Backend>: Clone + Send + core::fmt::Debug {
/// Type to save and load the module.
    type Record: Record<B>;
/// Return all the devices found in the underneath module tree added to the given vector
/// without duplicates.
    fn collect_devices(&self, devices: Devices<B>) -> Devices<B>;
/// Return all the devices found in the underneath module tree without duplicates.
    fn devices(&self) -> Devices<B> {
/// Fork the module and all of its sub-modules to the given device.
///
/// # Notes
///
/// This is similar to [to_device](Module::to_device), but it ensures the output module on the
/// new device will have its own autodiff graph.
    fn fork(self, device: &B::Device) -> Self;
/// Move the module and all of its sub-modules to the given device.
///
/// # Warnings
///
/// The operation supports autodiff and it will be registered when activated. However, this may
/// not be what you want. The output model will be an intermediary model, meaning that you
/// can't optimize it with gradient descent. If you want to optimize the output network on the
/// target device, use [fork](Module::fork) instead.
    fn to_device(self, device: &B::Device) -> Self;
/// Each tensor in the module tree will not require grad.
///
/// # Warnings
///
/// This should not be used for inference, use [valid](AutodiffModule::valid) when using
/// AD modules. This is mostly useful when performing partial finetuning, which is updating only
/// a small fraction of the parameters instead of finetuning all of them.
    fn no_grad(self) -> Self {
/// Get the number of parameters the module has, including all of its sub-modules.
    fn num_params(&self) -> usize {
/// Visit each tensor parameter in the module with a [visitor](ModuleVisitor).
    fn visit<Visitor: ModuleVisitor<B>>(&self, visitor: &mut Visitor);
/// Map each tensor parameter in the module with a [mapper](ModuleMapper).
    fn map<Mapper: ModuleMapper<B>>(self, mapper: &mut Mapper) -> Self;
/// Load the module state from a record.
    fn load_record(self, record: Self::Record) -> Self;
/// Convert the module into a record containing the state.
    fn into_record(self) -> Self::Record;
/// Save the module to a file using the provided [file recorder](crate::record::FileRecorder).
///
/// List of supported file recorders:
///
/// * [default](crate::record::DefaultFileRecorder)
/// * [bincode](crate::record::BinFileRecorder)
/// * [bincode compressed with gzip](crate::record::BinGzFileRecorder)
/// * [json pretty](crate::record::PrettyJsonFileRecorder)
/// * [json compressed with gzip](crate::record::JsonGzFileRecorder)
/// * [named mpk](crate::record::NamedMpkFileRecorder)
/// * [named mpk compressed with gzip](crate::record::NamedMpkGzFileRecorder)
///
/// ## Notes
///
/// The file extension is automatically added depending on the file recorder provided, you
/// don't have to specify it.
    fn save_file<FR, PB>(
        self,
        file_path: PB,
        recorder: &FR,
    ) -> Result<(), crate::record::RecorderError>
    where
        FR: crate::record::FileRecorder<B>,
        PB: Into<std::path::PathBuf>,
    {
/// Load the module from a file using the provided [file recorder](crate::record::FileRecorder).
///
/// The recorder should be the same as the one used to save the module, see
/// [save_file](Self::save_file).
///
/// ## Notes
///
/// The file extension is automatically added depending on the file recorder provided, you
/// don't have to specify it.
    fn load_file<FR, PB>(
        self,
        file_path: PB,
        recorder: &FR,
        device: &B::Device,
    ) -> Result<Self, crate::record::RecorderError>
    where
        FR: crate::record::FileRecorder<B>,
        PB: Into<std::path::PathBuf>,
    {
/// Quantize the weights of the module.
    fn quantize_weights(self, quantizer: &mut Quantizer) -> Self {
/// Module visitor trait for traversing and inspecting module parameters.
pub trait ModuleVisitor<B: Backend> {
/// Visit a float parameter in the module.
///
/// # Parameters
/// - `param`: The float parameter to visit
    fn visit_float<const D: usize>(&mut self, param: &Param<Tensor<B, D>>) {
/// Visit an int parameter in the module.
///
/// # Parameters
/// - `param`: The integer parameter to visit
    fn visit_int<const D: usize>(&mut self, param: &Param<Tensor<B, D, Int>>) {
/// Visit a bool parameter in the module.
///
/// # Parameters
/// - `param`: The boolean parameter to visit
    fn visit_bool<const D: usize>(&mut self, param: &Param<Tensor<B, D, Bool>>) {
/// Called when entering a submodule.
///
/// # Parameters
/// - `name`: The name of the submodule being entered
/// - `container_type`: The type of the container (e.g., "Module", "Vec", etc.)
    fn enter_module(&mut self, name: &str, container_type: &str) {
/// Called when exiting a submodule.
///
/// # Parameters
/// - `name`: The name of the submodule being exited
/// - `container_type`: The type of the container (e.g., "Module", "Vec", etc.)
    fn exit_module(&mut self, name: &str, container_type: &str) {
/// Visit a float tensor with its full module path.
///
/// # Parameters
/// - `path`: The path components to the tensor as a slice (e.g., &["encoder", "layer1", "weight"]).
///   Each element represents a module name in the hierarchy, with the final element
///   being the parameter name. This allows efficient reuse of the path stack.
/// - `id`: The unique identifier of the parameter
/// - `tensor`: The float tensor to visit
    fn visit_float_with_path<const D: usize>(
        &mut self,
        path: &[String],
        id: ParamId,
        tensor: &Tensor<B, D>,
    ) {
/// Visit an int tensor with its full module path.
///
/// # Parameters
/// - `path`: The path components to the tensor as a slice (e.g., &["encoder", "layer1", "weight"]).
///   Each element represents a module name in the hierarchy, with the final element
///   being the parameter name. This allows efficient reuse of the path stack.
/// - `id`: The unique identifier of the parameter
/// - `tensor`: The integer tensor to visit
    fn visit_int_with_path<const D: usize>(
        &mut self,
        path: &[String],
        id: ParamId,
        tensor: &Tensor<B, D, Int>,
    ) {
/// Visit a bool tensor with its full module path.
///
/// # Parameters
/// - `path`: The path components to the tensor as a slice (e.g., &["encoder", "layer1", "weight"]).
///   Each element represents a module name in the hierarchy, with the final element
///   being the parameter name. This allows efficient reuse of the path stack.
/// - `id`: The unique identifier of the parameter
/// - `tensor`: The boolean tensor to visit
    fn visit_bool_with_path<const D: usize>(
        &mut self,
        path: &[String],
        id: ParamId,
        tensor: &Tensor<B, D, Bool>,
    ) {
/// Module mapper trait for transforming module parameters.
pub trait ModuleMapper<B: Backend> {
/// Called when entering a submodule.
///
/// # Parameters
/// - `name`: The name of the submodule being entered
/// - `container_type`: The type of the container (e.g., "Module", "Vec", etc.)
    fn enter_module(&mut self, name: &str, container_type: &str) {
/// Called when exiting a submodule.
///
/// # Parameters
/// - `name`: The name of the submodule being exited
/// - `container_type`: The type of the container (e.g., "Module", "Vec", etc.)
    fn exit_module(&mut self, name: &str, container_type: &str) {
/// Map a float parameter in the module.
///
/// # Parameters
/// - `param`: The float parameter to transform
///
/// # Returns
/// The transformed parameter
    fn map_float<const D: usize>(&mut self, param: Param<Tensor<B, D>>) -> Param<Tensor<B, D>> {
/// Map an int parameter in the module.
///
/// # Parameters
/// - `param`: The integer parameter to transform
///
/// # Returns
/// The transformed parameter
    fn map_int<const D: usize>(
        &mut self,
        param: Param<Tensor<B, D, Int>>,
    ) -> Param<Tensor<B, D, Int>> {
/// Map a bool parameter in the module.
///
/// # Parameters
/// - `param`: The boolean parameter to transform
///
/// # Returns
/// The transformed parameter
    fn map_bool<const D: usize>(
        &mut self,
        param: Param<Tensor<B, D, Bool>>,
    ) -> Param<Tensor<B, D, Bool>> {
/// Module with auto-differentiation backend.
pub trait AutodiffModule<B: AutodiffBackend>: Module<B> + Send + core::fmt::Debug {
/// Inner module without auto-differentiation.
    type InnerModule: Module<B::InnerBackend>;
/// Get the same module, but on the inner backend without auto-differentiation.
    fn valid(&self) -> Self::InnerModule;
/// Default display settings for a module.
pub trait ModuleDisplayDefault {
/// Attributes of the module used for display purposes.
///
/// # Arguments
///
/// * `_content` - The content object that contains display settings and attributes.
///
/// # Returns
///
/// An optional content object containing the display attributes.
    fn content(&self, _content: Content) -> Option<Content>;
/// Gets the number of the parameters of the module.
    fn num_params(&self) -> usize {
/// Trait to implement custom display settings for a module.
///
/// In order to implement custom display settings for a module,
/// 1. Add #[module(custom_display)] attribute to the module struct after #[derive(Module)]
/// 2. Implement ModuleDisplay trait for the module
pub trait ModuleDisplay: ModuleDisplayDefault {
/// Formats the module with provided display settings.
///
/// # Arguments
///
/// * `passed_settings` - Display settings passed to the module.
///
/// # Returns
///
/// A string representation of the formatted module.
    fn format(&self, passed_settings: DisplaySettings) -> String {
// Use custom content if it is implemented and show_all_attributes is false,
// otherwise use default content
// If there is only one item in the content, return it or no attributes
// Print the struct name
/// Custom display settings for the module.
///
/// # Returns
///
/// An optional display settings object.
    fn custom_settings(&self) -> Option<DisplaySettings> {
/// Custom attributes for the module.
///
/// # Arguments
///
/// * `_content` - The content object that contains display settings and attributes.
///
/// # Returns
///
/// An optional content object containing the custom attributes.
    fn custom_content(&self, _content: Content) -> Option<Content> {
/// Custom module display settings.
pub struct DisplaySettings {
/// Whether to print the module parameter ids.
/// Whether to print the module attributes.
/// Whether to print the module number of parameters.
/// Print new line after an attribute.
/// Indentation size.
/// Level of indentation.
impl Default for DisplaySettings {
    fn default() -> Self {
impl DisplaySettings {
/// Create a new format settings.
///
/// # Returns
///
/// A new instance of `DisplaySettings`.
    pub fn new() -> Self {
/// Sets a flag to show module parameters.
///
/// # Arguments
///
/// * `flag` - Boolean flag to show module parameters.
///
/// # Returns
///
/// Updated `DisplaySettings` instance.
    pub fn with_show_param_id(mut self, flag: bool) -> Self {
/// Sets a flag to show module attributes.
///
/// # Arguments
///
/// * `flag` - Boolean flag to show all module attributes.
///
/// # Returns
///
/// Updated `DisplaySettings` instance.
    pub fn with_show_all_attributes(mut self, flag: bool) -> Self {
/// Sets a flag to show the number of module parameters.
///
/// # Arguments
///
/// * `flag` - Boolean flag to show the number of module parameters.
///
/// # Returns
///
/// Updated `DisplaySettings` instance.
    pub fn with_show_num_parameters(mut self, flag: bool) -> Self {
/// Sets a flag to print a new line after an attribute.
///
/// # Arguments
///
/// * `flag` - Boolean flag to print a new line after an attribute.
///
/// # Returns
///
/// Updated `DisplaySettings` instance.
    pub fn with_new_line_after_attribute(mut self, flag: bool) -> Self {
/// Sets the indentation size.
///
/// # Arguments
///
/// * `size` - The size of the indentation.
///
/// # Returns
///
/// Updated `DisplaySettings` instance.
    pub fn with_indentation_size(mut self, size: usize) -> Self {
/// Inherits settings from the provided settings and return a new settings object.
///
/// # Arguments
///
/// * `top` - The top level `DisplaySettings` to inherit from.
///
/// # Returns
///
/// Updated `DisplaySettings` instance.
    pub fn inherit(self, top: Self) -> Self {
/// A convenience method to wrap the DisplaySettings struct in an option.
///
/// # Returns
///
/// An optional `DisplaySettings`.
    pub fn optional(self) -> Option<Self> {
/// Increases the level of indentation.
///
/// # Returns
///
/// Updated `DisplaySettings` instance with increased indentation level.
    pub fn level_up(mut self) -> Self {
/// Gets `show_param_id` flag, substitutes false if not set.
///
/// This flag is used to print the module parameter ids.
///
/// # Returns
///
/// A boolean value indicating whether to show parameter ids.
    pub fn show_param_id(&self) -> bool {
/// Gets `show_all_attributes`, substitutes false if not set.
///
/// This flag is used to force to print all module attributes, overriding custom attributes.
///
/// # Returns
///
/// A boolean value indicating whether to show all attributes.
    pub fn show_all_attributes(&self) -> bool {
/// Gets `show_num_parameters`, substitutes true if not set.
///
/// This flag is used to print the number of module parameters.
///
/// # Returns
///
/// A boolean value indicating whether to show the number of parameters.
    pub fn show_num_parameters(&self) -> bool {
/// Gets `new_line_after_attribute`, substitutes true if not set.
///
/// This flag is used to print a new line after an attribute.
///
/// # Returns
///
/// A boolean value indicating whether to print a new line after an attribute.
    pub fn new_line_after_attribute(&self) -> bool {
/// Gets `indentation_size`, substitutes 2 if not set.
///
/// This flag is used to set the size of indentation.
///
/// # Returns
///
/// An integer value indicating the size of indentation.
    pub fn indentation_size(&self) -> usize {
/// Struct to store the attributes of a module for formatting.
pub struct Content {
/// List of attributes.
/// Single item content.
/// Display settings.
/// Top level type name.
impl Content {
/// Creates a new attributes struct.
///
/// # Arguments
///
/// * `display_settings` - Display settings for the content.
///
/// # Returns
///
/// A new instance of `Content`.
    pub fn new(display_settings: DisplaySettings) -> Self {
/// Adds an attribute to the format settings. The value will be formatted and stored as a string.
///
/// # Arguments
///
/// * `name` - Name of the attribute.
/// * `value` - Value of the attribute.
///
/// # Returns
///
/// Updated `Content` instance with the new attribute added.
    pub fn add<T: ModuleDisplay + ?Sized>(mut self, name: &str, value: &T) -> Self {
// TODO level + 1
/// Adds a single item.
///
/// # Arguments
///
/// * `value` - Rendered string of the single item.
///
/// # Returns
///
/// Updated `Content` instance with the single item added.
    pub fn add_single<T: ModuleDisplay + ?Sized>(mut self, value: &T) -> Self {
/// Adds a single item.
///
/// # Arguments
///
/// * `value` - Formatted display value.
///
/// # Returns
///
/// Updated `Content` instance with the formatted single item added.
    pub fn add_formatted<T: Display>(mut self, value: &T) -> Self {
/// A convenience method to wrap the Attributes struct in an option
/// because it is often used as an optional field.
///
/// # Returns
///
/// An optional `Content`.
    pub fn optional(self) -> Option<Self> {
/// Sets the top level type name.
///
/// # Arguments
///
/// * `ty` - The type name to set.
///
/// # Returns
///
/// Updated `Content` instance with the top level type name set.
    pub fn set_top_level_type(mut self, ty: &str) -> Self {
/// Attribute to print in the display method.
pub struct Attribute {
/// Name of the attribute.
/// Value of the attribute.
/// Type of the attribute.
/// Extracts the short name of a type T
///
/// # Returns
///
/// A string slice representing the short name of the type.
pub fn extract_type_name<T: ?Sized>() -> &'static str {
// Get the full type name of T, including module path and generic parameters
// Find the first occurrence of '<' in the full type name
// If not found, use the length of the type name
// Slice the type name up to the first '<' or the end
// Find the last occurrence of "::" in the sliced type name
// If found, add 2 to skip the "::" itself
// If not found, start from the beginning of the type name
// Find the last occurrence of '<' in the sliced type name
// If not found, use the length of the type name
// If the start index is less than the end index,
// return the slice of the type name from start to end
// Otherwise, return the entire sliced type name
/// Enum specifying with what values a tensor should be initialized
pub enum Initializer {
/// Fills tensor with specified value everywhere
/// The value to fill the tensor with
/// Fills tensor with 1s everywhere
/// Fills tensor with 0s everywhere
/// Fills tensor with values drawn uniformly between specified values
/// The minimum value to draw from
/// The maximum value to draw from
/// Fills tensor with values drawn from normal distribution with specified mean and std
/// The mean of the normal distribution
/// The standard deviation of the normal distribution
/// Fills tensor with values according to the uniform version of Kaiming initialization
/// The gain to use in initialization formula
/// Whether to use fan out only in initialization formula
/// Fills tensor with values according to the uniform version of Kaiming initialization
/// The gain to use in initialization formula
/// Whether to use fan out only in initialization formula
/// Fills tensor with values according to the uniform version of Xavier Glorot initialization
/// described in [Understanding the difficulty of training deep feedforward neural networks
/// ](https://proceedings.mlr.press/v9/glorot10a/glorot10a.pdf)
/// The gain to use in initialization formula
/// Fills tensor with values according to the normal version of Xavier Glorot initialization
/// described in [Understanding the difficulty of training deep feedforward neural networks
/// ](https://proceedings.mlr.press/v9/glorot10a/glorot10a.pdf)
/// The gain to use in initialization formula
/// Fills tensor with values according to the (semi) orthogonal initialization
/// described in [Exact solutions to the nonlinear dynamics of learning in deep linear neural networks`
///  - [Saxe, A. et al. (2013)](https://arxiv.org/abs/1312.6120)
/// The gain to use in initialization formula
impl Initializer {
/// Inits a tensor parameter of given shape with values depending on initializer kind.
///
/// # Params
///
/// - shape: Shape of the initiated tensor.
    pub fn init<B: Backend, const D: usize, S: Into<Shape>>(
        &self,
        shape: S,
        device: &B::Device,
    ) -> Param<Tensor<B, D>> {
/// Inits a tensor parameter of given shape with values depending on initializer kind.
///
/// # Params
///
/// - shape: Shape of the initiated tensor.
    pub fn init_with<B: Backend, const D: usize, S: Into<Shape>>(
        &self,
        shape: S,
        fan_in: Option<usize>,
        fan_out: Option<usize>,
        device: &B::Device,
    ) -> Param<Tensor<B, D>> {
    fn init_tensor<B: Backend, const D: usize, S: Into<Shape>>(
        &self,
        shape: S,
        fan_in: Option<usize>,
        fan_out: Option<usize>,
        device: &B::Device,
    ) -> Tensor<B, D> {
// following the implementation in pytorch:
// https://github.com/pytorch/pytorch/blob/v2.7.0/torch/nn/init.py#L574
    fn kaiming_std(
        &self,
        fan_out_only: bool,
        fan_in: Option<usize>,
        fan_out: Option<usize>,
    ) -> f64 {
    fn xavier_std(&self, fan_in: Option<usize>, fan_out: Option<usize>) -> f64 {
fn uniform_draw<B: Backend, const D: usize, S: Into<Shape>>(
    shape: S,
    low: f64,
    high: f64,
    device: &B::Device,
) -> Tensor<B, D> {
fn normal_draw<B: Backend, const D: usize, S: Into<Shape>>(
    shape: S,
    mean: f64,
    std: f64,
    device: &B::Device,
) -> Tensor<B, D> {
fn qr_decomposition<B: Backend>(
    a: Tensor<B, 2>,
    device: &B::Device,
) -> (Tensor<B, 2>, Tensor<B, 2>) {
// Calculate the QR decomposition using Gram-Schmidt-process: https://en.wikipedia.org/wiki/Gram%E2%80%93Schmidt_process
// norm of v
    pub type TB = burn_ndarray::NdArray<f32>;
    type FT = FloatElem<TB>;
    fn assert_normal_init(expected_mean: f64, expected_var: f64, tensor: &Tensor<TB, 2>) {
    fn initializer_uniform_init() {
    fn initializer_normal_init() {
// seed random generator
    fn initializer_constant_init() {
    fn initializer_zeros_init() {
    fn initializer_ones_init() {
    fn initializer_kaiming_uniform_init() {
    fn initializer_kaiming_normal_init() {
    fn initializer_kaiming_uniform_init_bias() {
    fn initializer_kaiming_uniform_init_fan_out() {
    fn initializer_kaiming_uniform_no_fan() {
    fn initializer_xavier_uniform_init() {
    fn initializer_xavier_normal_init() {
    fn initializer_xavier_uniform_no_fan() {
    fn test_qr_decomposition() {
// test values follow the example from https://pytorch.org/docs/stable/generated/torch.linalg.qr.html#torch.linalg.qr
// Q @ R should reconstruct input `a`
// assert that the difference between input (`a`) and Q @ R is (almost) zero
    fn initializer_orthogonal_correct() {
// test 2D tensor
// Q.T @ Q should be close to identity matrix
    fn initializer_orthogonal_init() {
// test 2D tensor
// test 3D tensor
    fn initializer_orthogonal_init_1d() {
// test 1D tensor
type Mapper<T> = Arc<dyn Fn(T) -> T + Send + Sync>;
type Mapper<T> = Arc<Box<dyn Fn(T) -> T + Send + Sync>>;
fn new_mapper<T, F: Fn(T) -> T + Send + Sync + 'static>(func: F) -> Mapper<T> {
fn new_mapper<T, F: Fn(T) -> T + Send + Sync + 'static>(func: F) -> Mapper<T> {
/// Parameters are the fundamental building blocks of [modules](crate::module::Module) where they
/// serve as containers for [tensors](crate::tensor::Tensor) that can be updated during
/// training, and loaded during inference. If you don't want to save the tensors
/// and/or don't want to update it during training, you don't need this type to wrap your tensor.
///
/// # Core Lazy Initialization Architecture
///
/// `Param<T>` has a dual-state design using `OnceCell<T>`:
///
/// ## State Management
///
/// **Two possible states:**
///
/// 1. **Initialized**: `state: OnceCell<T>` contains value, `initialization: None`
/// 2. **Uninitialized (Lazy)**: `state` is empty, `initialization: Some(RwLock<Option<Uninitialized<T>>>)`
pub struct Param<T: Parameter> {
/// The unique ID of this parameter. This is used by eg. optimizers to associate a gradient with a specific parameter.
/// The OnceCell holding the initialized parameter value.
/// Empty for uninitialized parameters, populated after first access or explicit initialization.
/// The deferred initialization state for lazy parameters.
///
/// **State Transitions:**
/// - Initialized params: `None`
/// - Uninitialized params: `Some(RwLock<Some(Uninitialized<T>)>)`
/// - After lazy init triggers: `Some(RwLock<None>)` (inner Option is taken)
/// Applies transformations when loading and saving parameters.
///
/// # Mapper System
///
/// `ParamMapper<T>` allows applying transformations during serialization and deserialization:
/// - `load: Option<Mapper<T>>` - transformation during deserialization (applied in `transform_for_load()`)
/// - `save: Option<Mapper<T>>` - transformation during serialization (applied in `transform_for_save()`)
///
/// These are commonly used for:
/// - Quantization/dequantization
/// - Precision conversion (e.g., FP32 ↔ FP16)
/// - Custom parameter transformations
pub struct ParamMapper<T: Parameter> {
impl<T: Parameter> core::fmt::Debug for ParamMapper<T> {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
impl<T: Parameter> ParamMapper<T> {
/// Applies the transformation when loading the given parameter.
    pub fn on_load(&self, param: T) -> T {
/// Applies the transformation when saving the given parameter.
    pub fn on_save(&self, param: T) -> T {
impl<T: Parameter> Default for ParamMapper<T> {
    fn default() -> Self {
impl<T: Parameter> core::fmt::Display for Param<T> {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
impl<T: Parameter> core::fmt::Debug for Param<T> {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
/// Trait that defines what is necessary for a type to be a parameter.
pub trait Parameter: Clone + core::fmt::Debug + Send {
/// The device type to be used.
    type Device: Clone;
/// Fetch the device.
    fn device(&self) -> Self::Device;
/// Fetch the gradient requirement.
    fn is_require_grad(&self) -> bool;
/// Set the gradient requirement.
    fn set_require_grad(self, require_grad: bool) -> Self;
/// The deferred initialization state for lazy parameters.
pub(crate) struct Uninitialized<P: Parameter> {
/// The initialization function. Called with `(device, is_require_grad) -> Parameter`.
/// This function is consumed during initialization via `FnOnce`.
/// The target device on which the parameter should be initialized.
/// Used by `lazy_device()` to provide device information without triggering initialization.
/// The gradient requirement for the parameter.
/// Used by `lazy_is_require_grad()` to provide gradient settings without triggering initialization.
/// The shape of the tensor parameter.
/// Used by `lazy_shape()` to provide shape information without triggering initialization.
impl<P: Parameter> Uninitialized<P> {
/// Consumes the uninitialized state and runs the initialization function.
///
/// This is called by [Param::val] when accessing an uninitialized parameter for the first time.
/// The function is given the stored device and gradient requirement, and returns the initialized parameter.
    fn initialize(self) -> P {
impl<T: Parameter> Param<T> {
/// Create a new parameter that is already initialized.
    pub fn initialized(id: ParamId, value: T) -> Self {
/// Create a new parameter that is not already initialized.
    pub fn uninitialized<F>(
        id: ParamId,
        init: F,
        device: T::Device,
        is_require_grad: bool,
        shape: Shape,
    ) -> Self
    where
        F: FnOnce(&T::Device, bool) -> T + Send + 'static,
    {
/// Gets the parameter value, initializing it lazily if needed.
///
/// For initialized parameters, this returns a clone of the cached value.
/// For uninitialized parameters, this triggers initialization:
    pub fn val(&self) -> T {
/// Check if the parameter has been initialized.
///
/// Returns `true` if the parameter's value has been computed and cached,
/// `false` if it's still lazy and will be initialized on first access.
    pub fn is_initialized(&self) -> bool {
/// Gets the parameter's value while consuming the parameter.
    pub fn into_value(self) -> T {
/// Gets the parameter id and value while consuming the parameter.
    pub fn consume(self) -> (ParamId, T, ParamMapper<T>) {
/// Execute the given function on the inner value.
    pub fn map<F: FnOnce(T) -> T>(self, func: F) -> Self {
/// Create an initialized parameter with the given id, value, and param mapper.
///
/// This is a helper method for creating parameters while preserving the param mapper,
/// typically used in ModuleMapper implementations.
    pub fn from_mapped_value(id: ParamId, value: T, param_mapper: ParamMapper<T>) -> Self {
/// Runs a transformation on the parameter when loading.
    pub fn load_mapper<F: Fn(T) -> T + Send + Sync + 'static>(mut self, func: F) -> Self {
/// Runs a transformation on the parameter when saving.
    pub fn save_mapper<F: Fn(T) -> T + Send + Sync + 'static>(mut self, func: F) -> Self {
/// Execute the given function on the inner value.
    pub fn init_mapper<F: FnOnce(T) -> T + Send + 'static>(self, func: F) -> Self
    where
        T: 'static,
    {
/// The device on which the parameter is or will be initialized, **without triggering initialization**.
///
/// This is critical for the load optimization: when loading tensors into an uninitialized parameter,
/// we need to know the target device to move the loaded tensor appropriately, but we don't want to
/// trigger the initialization function (which would allocate an unnecessary tensor).
///
/// Use this instead of [crate::tensor::Tensor::device] when you need the device but want to
/// preserve lazy initialization.
    pub fn lazy_device(&self) -> T::Device {
/// The gradient requirement on which the parameter is or will be initialized, **without triggering initialization**.
///
/// Similar to [lazy_device](Self::lazy_device), this is critical for the load optimization.
/// When loading tensors into an uninitialized parameter, we need to apply the correct gradient
/// setting to the loaded tensor without triggering the initialization function.
///
/// # Notes
///
/// This is a crate-private function, since users are not expected to use `is_require_grad` of an
/// uninitialized module to then override its value. All low-level functions should be provided
/// by `burn` and should handle those details.
    pub(crate) fn lazy_is_require_grad(&self) -> bool {
/// Override the gradient requirement for the current parameter.
    pub fn set_require_grad(self, require_grad: bool) -> Self {
impl<T: Parameter> Clone for Param<T> {
    fn clone(&self) -> Self {
impl<T: Parameter> Deref for Param<T> {
    type Target = T;
    fn deref(&self) -> &Self::Target {
/// Record used for constant type implementing the [module](crate::module::Module) trait.
pub struct ConstantRecord;
impl serde::Serialize for ConstantRecord {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: serde::Serializer,
    {
// nothing to serialize
impl<'de> serde::Deserialize<'de> for ConstantRecord {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
impl<B: Backend> Record<B> for ConstantRecord {
    type Item<S: PrecisionSettings> = ConstantRecord;
    fn into_item<S: PrecisionSettings>(self) -> Self::Item<S> {
    fn from_item<S: PrecisionSettings>(item: Self::Item<S>, _device: &B::Device) -> Self {
/// Constant macro.
        type Record = burn::module::ConstantRecord;
        fn visit<V: burn::module::ModuleVisitor<B>>(&self, _visitor: &mut V) {
// Nothing to do
        fn map<M: burn::module::ModuleMapper<B>>(self, _mapper: &mut M) -> Self {
        fn load_record(self, _record: Self::Record) -> Self {
        fn into_record(self) -> Self::Record {
        fn to_device(self, _: &B::Device) -> Self {
        fn fork(self, _: &B::Device) -> Self {
        fn collect_devices(&self, devices: burn::module::Devices<B>) -> burn::module::Devices<B> {
    (ad_module, $type:ty) => {
        type InnerModule = $type;
        fn valid(&self) -> Self::InnerModule {
    ($type:ty) => {
        impl<B: burn::tensor::backend::Backend> burn::module::Module<B> for $type {
        impl<B: burn::tensor::backend::AutodiffBackend> burn::module::AutodiffModule<B> for $type {
            constant!(ad_module, $type);
        impl burn::module::ModuleDisplayDefault for $type {
            fn content(&self, content: burn::module::Content) -> Option<burn::module::Content> {
        impl burn::module::ModuleDisplay for $type {
// General Types
// Float Types
// Unsigned Integer Types
// Signed Integer Types
impl burn::module::ModuleDisplay for str {
impl burn::module::ModuleDisplayDefault for str {
    fn content(&self, content: burn::module::Content) -> Option<burn::module::Content> {
impl<const D: usize, B: Backend, K: BasicOps<B>> Module<B> for Tensor<B, D, K> {
    type Record = ConstantRecord;
    fn visit<V: ModuleVisitor<B>>(&self, _visitor: &mut V) {
    fn map<M: ModuleMapper<B>>(self, _mapper: &mut M) -> Self {
    fn into_record(self) -> Self::Record {
    fn load_record(self, _record: Self::Record) -> Self {
    fn to_device(self, device: &B::Device) -> Self {
    fn fork(self, device: &B::Device) -> Self {
    fn collect_devices(&self, mut devices: Devices<B>) -> Devices<B> {
impl<const D: usize, B: Backend, K: BasicOps<B>> ModuleDisplayDefault for Tensor<B, D, K> {
    fn content(&self, content: Content) -> Option<Content> {
impl<const D: usize, B: Backend, K: BasicOps<B>> ModuleDisplay for Tensor<B, D, K> {
impl<const D: usize, B: AutodiffBackend, K: BasicAutodiffOps<B>> AutodiffModule<B>
    for Tensor<B, D, K>
{
    type InnerModule = Tensor<B::InnerBackend, D, K::InnerKind>;
    fn valid(&self) -> Self::InnerModule {
impl<B: Backend> Module<B> for PhantomData<B> {
    type Record = ConstantRecord;
    fn visit<V: ModuleVisitor<B>>(&self, _visitor: &mut V) {
// Nothing to do
    fn map<M: ModuleMapper<B>>(self, _mapper: &mut M) -> Self {
    fn load_record(self, _record: Self::Record) -> Self {
    fn into_record(self) -> Self::Record {
    fn to_device(self, _: &Device<B>) -> Self {
    fn fork(self, _: &Device<B>) -> Self {
    fn collect_devices(&self, devices: Devices<B>) -> Devices<B> {
impl<B: Backend> ModuleDisplayDefault for PhantomData<B> {
    fn content(&self, content: Content) -> Option<Content> {
impl<B: Backend> ModuleDisplay for PhantomData<B> {
impl<B: AutodiffBackend> AutodiffModule<B> for PhantomData<B> {
    type InnerModule = PhantomData<B::InnerBackend>;
    fn valid(&self) -> Self::InnerModule {
/// Container to satisfy the Module trait for types that are not modules.
pub struct Ignored<T>(pub T);
impl<B, T> Module<B> for Ignored<T>
where
    B: Backend,
    T: Sync + Send + core::fmt::Debug + Clone,
{
    type Record = ConstantRecord;
    fn visit<V: ModuleVisitor<B>>(&self, _visitor: &mut V) {
// Nothing to do
    fn map<M: ModuleMapper<B>>(self, _mapper: &mut M) -> Self {
    fn load_record(self, _record: Self::Record) -> Self {
    fn into_record(self) -> Self::Record {
    fn to_device(self, _: &Device<B>) -> Self {
    fn fork(self, _: &Device<B>) -> Self {
    fn collect_devices(&self, devices: Devices<B>) -> Devices<B> {
impl<T> ModuleDisplayDefault for Ignored<T>
where
    T: Sync + Send + core::fmt::Debug + Clone,
{
    fn content(&self, content: Content) -> Option<Content> {
// For now, just print the debug representation of the ignored value
impl<T> ModuleDisplay for Ignored<T> where T: Sync + Send + core::fmt::Debug + Clone {
impl<T> Display for Ignored<T>
where
    T: Sync + Send + core::fmt::Debug + Clone,
{
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
impl<B: AutodiffBackend, T> AutodiffModule<B> for Ignored<T>
where
    B: AutodiffBackend,
    T: Sync + Send + core::fmt::Debug + Clone,
{
    type InnerModule = Ignored<T>;
    fn valid(&self) -> Self::InnerModule {
// Implement deref for Ignored
impl<T> core::ops::Deref for Ignored<T> {
    type Target = T;
    fn deref(&self) -> &Self::Target {
    fn tensor_load_record_setting() {
    fn empty_module_with_phantom() {
        struct EmptyModule<B: Backend> {
// Hashbrown changed its default hasher in 0.15, but there are some issues
// https://github.com/rust-lang/hashbrown/issues/577
// Also, `param_serde_deserialize_legacy_uuid` doesn't pass with the default hasher.
type DefaultHashBuilder = core::hash::BuildHasherDefault<ahash::AHasher>;
/// Parameter ID.
pub struct ParamId {
impl From<u64> for ParamId {
    fn from(value: u64) -> Self {
impl Default for ParamId {
    fn default() -> Self {
impl ParamId {
/// Create a new parameter ID.
    pub fn new() -> Self {
/// Gets the internal value of the id.
    pub fn val(&self) -> u64 {
/// Convert the parameter ID into a string.
    pub fn serialize(self) -> String {
/// Deserialize a param id.
///
/// Preserves compatibility with previous formats (6 bytes, 16-byte uuid).
    pub fn deserialize(encoded: &str) -> ParamId {
// Backward compatibility with uuid parameter identifiers
// Hash the 128-bit uuid to 64-bit
// Though not *theoretically* unique, the probability of a collision should be extremely low
// let mut hasher = DefaultHasher::new();
impl core::fmt::Display for ParamId {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
    fn param_serde_deserialize() {
    fn param_serde_deserialize_legacy() {
    fn param_serde_deserialize_legacy_uuid() {
// Ensure support for legacy uuid deserialization and make sure it results in the same output
    fn param_serde_deserialize_invalid_id() {
impl<T, B> Module<B> for Option<T>
where
    T: Module<B> + Debug + Send + Clone,
    B: Backend,
{
    type Record = Option<T::Record>;
    fn visit<V: ModuleVisitor<B>>(&self, visitor: &mut V) {
    fn map<M: ModuleMapper<B>>(self, mapper: &mut M) -> Self {
    fn load_record(self, record: Self::Record) -> Self {
    fn into_record(self) -> Self::Record {
    fn to_device(self, device: &Device<B>) -> Self {
    fn fork(self, device: &Device<B>) -> Self {
    fn collect_devices(&self, mut devices: Vec<B::Device>) -> Vec<B::Device> {
impl<T: ModuleDisplay> ModuleDisplayDefault for Option<T> {
    fn content(&self, content: Content) -> Option<Content> {
impl<T: ModuleDisplay> ModuleDisplay for Option<T> {
impl<T, B> AutodiffModule<B> for Option<T>
where
    T: AutodiffModule<B> + Debug + Send + Clone,
    B: AutodiffBackend,
{
    type InnerModule = Option<T::InnerModule>;
    fn valid(&self) -> Self::InnerModule {
impl<T, B> Module<B> for Vec<T>
where
    T: Module<B> + Debug + Send + Clone,
    B: Backend,
{
    type Record = Vec<T::Record>;
    fn num_params(&self) -> usize {
    fn visit<V: ModuleVisitor<B>>(&self, visitor: &mut V) {
    fn map<M: ModuleMapper<B>>(self, mapper: &mut M) -> Self {
    fn into_record(self) -> Self::Record {
    fn load_record(self, record: Self::Record) -> Self {
    fn to_device(self, device: &Device<B>) -> Self {
    fn fork(self, device: &Device<B>) -> Self {
    fn collect_devices(&self, mut devices: Vec<B::Device>) -> Vec<B::Device> {
impl<T: ModuleDisplay> ModuleDisplayDefault for Vec<T> {
    fn content(&self, content: Content) -> Option<Content> {
impl<T: ModuleDisplay> ModuleDisplay for Vec<T> {
impl<T, B> AutodiffModule<B> for Vec<T>
where
    T: AutodiffModule<B> + Debug + Send + Clone,
    B: AutodiffBackend,
{
    type InnerModule = Vec<T::InnerModule>;
    fn valid(&self) -> Self::InnerModule {
impl<const N: usize, T, B> Module<B> for [T;
    type Record = [T::Record;
    fn collect_devices(&self, mut devices: Vec<B::Device>) -> Vec<B::Device> {
    fn num_params(&self) -> usize {
    fn visit<V: ModuleVisitor<B>>(&self, visitor: &mut V) {
    fn map<M: ModuleMapper<B>>(self, mapper: &mut M) -> Self {
    fn load_record(self, record: Self::Record) -> Self {
    fn into_record(self) -> Self::Record {
    fn to_device(self, device: &Device<B>) -> Self {
    fn fork(self, device: &Device<B>) -> Self {
impl<const N: usize, T: ModuleDisplay> ModuleDisplayDefault for [T;
    fn content(&self, content: Content) -> Option<Content> {
impl<const N: usize, T: ModuleDisplay> ModuleDisplay for [T;
impl<const N: usize, T, B> AutodiffModule<B> for [T;
    type InnerModule = [T::InnerModule;
    fn valid(&self) -> Self::InnerModule {
/// A macro for generating implementations for tuple modules of different sizes.
/// For example: `impl_module_tuple!([L0, L1][0, 1])`.
/// Would generate an implementation for a tuple of size 2.
/// For this macro to work properly, please adhere to the convention:
/// `impl_module_tuple!([L0, L1, ..., Ln][0, 1, ..., n])`.
// `$l` represents the generic modules.
// `$i` represents the indices of the modules in the tuple.
        impl<B, $($l,)*> Module<B> for ($($l,)*)
        where
            B: Backend,
            $($l: Module<B> + Debug + Send + Clone,)*
        {
            type Record = ($($l::Record),*);
            fn collect_devices(&self, mut devices: Vec<B::Device>) -> Vec<B::Device> {
            fn fork(self, device: &Device<B>) -> Self {
            fn to_device(self, device: &Device<B>) -> Self {
            fn visit<V: ModuleVisitor<B>>(&self, visitor: &mut V) {
            fn map<M: ModuleMapper<B>>(self, mapper: &mut M) -> Self {
            fn load_record(self, record: Self::Record) -> Self {
            fn into_record(self) -> Self::Record {
        impl<B, $($l,)*> AutodiffModule<B> for ($($l,)*)
        where
            B: AutodiffBackend,
            $($l: AutodiffModule<B> + Debug + Send + Clone,)*
        {
            type InnerModule = ($($l::InnerModule,)*);
            fn valid(&self) -> Self::InnerModule {
        impl<$($l,)*> ModuleDisplayDefault for ($($l,)*)
        where
            $($l: ModuleDisplay,)*
        {
            fn content(&self, content: Content) -> Option<Content> {
        impl<$($l,)*> ModuleDisplay for ($($l,)*) where $($l: ModuleDisplay,)* {
    fn dont_override_constant_module_when_loading_record() {
    fn dont_override_constant_module_when_loading_none_record() {
    pub(super) fn get_thread_current_id() -> ThreadId {
    pub(super) fn get_thread_current_id() -> ThreadId {
// Re-export items from the disabled/enabled blocks
/// A state that can be updated during the forward pass while being thread safe.
///
/// # Note
///
/// The state value is the average of all updates on all threads.
pub struct RunningState<V> {
// Implement display for the module
impl<V> core::fmt::Display for RunningState<V> {
    fn fmt(&self, f: &mut core::fmt::Formatter) -> core::fmt::Result {
impl<V> ModuleDisplayDefault for RunningState<V> {
    fn content(&self, content: Content) -> Option<Content> {
impl<V> ModuleDisplay for RunningState<V> {
impl<const D: usize, B: Backend> Module<B> for RunningState<Tensor<B, D>> {
    type Record = Param<Tensor<B, D>>;
    fn visit<V: ModuleVisitor<B>>(&self, visitor: &mut V) {
    fn map<M: ModuleMapper<B>>(self, mapper: &mut M) -> Self {
    fn into_record(self) -> Self::Record {
    fn load_record(mut self, record: Self::Record) -> Self {
    fn to_device(self, device: &Device<B>) -> Self {
    fn fork(self, device: &Device<B>) -> Self {
// Same thing here since no grad.
    fn collect_devices(&self, mut devices: Vec<Device<B>>) -> Vec<Device<B>> {
impl<const D: usize, B: Backend> RunningState<Tensor<B, D>> {
/// Create a new running state.
    pub fn new(value: Tensor<B, D>) -> Self {
/// Create a new running state.
    pub fn with_id(id: ParamId, value: Tensor<B, D>) -> Self {
/// Create a new running state from a record.
    pub fn from_record(record: Param<Tensor<B, D>>) -> Self {
/// Update the value on the current thread.
    pub fn update(&self, value: Tensor<B, D>) {
/// Get the current value,
///
/// # Note
///
/// The current value might be outdated by one update.
    pub fn value(&self) -> Tensor<B, D> {
/// Get the current value and make sure it is sync.
///
/// # Note
///
/// Don't use this function after an update on the same thread where other threads might have to
/// register their update before the actual synchronization needs to happen.
    pub fn value_sync(&self) -> Tensor<B, D> {
    fn sync(&self) {
    fn update_value(&self, map: &mut HashMap<ThreadId, Tensor<B, D>>) {
impl<const D: usize, B: AutodiffBackend> AutodiffModule<B> for RunningState<Tensor<B, D>> {
    type InnerModule = RunningState<Tensor<B::InnerBackend, D>>;
    fn valid(&self) -> Self::InnerModule {
impl<B: Backend, const D: usize> Parameter for Tensor<B, D, Float> {
    type Device = B::Device;
    fn device(&self) -> Self::Device {
    fn is_require_grad(&self) -> bool {
    fn set_require_grad(self, require_grad: bool) -> Self {
impl<B: Backend, const D: usize> Parameter for Tensor<B, D, Int> {
    type Device = B::Device;
    fn device(&self) -> Self::Device {
    fn is_require_grad(&self) -> bool {
    fn set_require_grad(self, _require_grad: bool) -> Self {
impl<B: Backend, const D: usize> Parameter for Tensor<B, D, Bool> {
    type Device = B::Device;
    fn device(&self) -> Self::Device {
    fn is_require_grad(&self) -> bool {
    fn set_require_grad(self, _require_grad: bool) -> Self {
impl<B: Backend, const D: usize> Param<Tensor<B, D>> {
/// Create a new parameter from a float tensor.
///
/// # Warnings
///
/// We strongly recommend using [Param::uninitialized] if you are using this method to
/// initialize parameters inside a module, since the tensor initialization will be lazy,
/// making the loading of weights more performant.
    pub fn from_tensor(value: Tensor<B, D>) -> Self {
// When creating a parameter from a float tensor, we automatically mark it as requiring
// gradients, so that it can be updated by an optimizer.
/// The shape of the parameter, **without triggering initialization**.
///
/// This is critical for shape validation during loading: when applying tensors to an
/// uninitialized parameter, we need to validate the shape without triggering the
/// initialization function (which would allocate an unnecessary tensor).
///
/// Use this instead of [crate::tensor::Tensor::shape] when you need the shape but want to
/// preserve lazy initialization.
    pub fn lazy_shape(&self) -> burn_tensor::Shape {
/// Create a new parameter from data.
    pub fn from_data<T>(data: T, device: &B::Device) -> Self
    where
        T: Into<TensorData>,
    {
// When creating a parameter from a float tensor, we automatically mark it as requiring
// gradients, so that it can be updated by an optimizer.
/// Transform a parameter for loading by applying load transformations.
///
/// This method is used to restore a parameter from a tensor (typically during deserialization).
/// It ensures the tensor is moved to the expected device, applies the param mapper's
/// `on_load` transformation, and preserves the autodiff settings (require_grad).
    pub fn transform_for_load(self, tensor: Tensor<B, D>, param_id: ParamId) -> Self {
// Make sure we load the tensor into the same module device.
// Make sure we load the tensor with the same autodiff setting.
/// Transform a parameter for saving by applying save transformations.
///
/// This method is used to prepare a parameter for saving (typically during serialization).
/// It applies the param mapper's `on_save` transformation, which can be used
/// to modify the tensor before serialization (e.g., quantization, precision conversion).
    pub fn transform_for_save(&self) -> Self {
impl<B: Backend, const D: usize> Param<Tensor<B, D, Int>> {
/// The shape of the parameter, **without triggering initialization**.
///
/// This is critical for shape validation during loading: when applying tensors to an
/// uninitialized parameter, we need to validate the shape without triggering the
/// initialization function (which would allocate an unnecessary tensor).
///
/// Use this instead of [crate::tensor::Tensor::shape] when you need the shape but want to
/// preserve lazy initialization.
    pub fn lazy_shape(&self) -> burn_tensor::Shape {
/// Transform a parameter for loading by applying load transformations.
///
/// This method is used to restore a parameter from a tensor (typically during deserialization).
/// It ensures the tensor is moved to the expected device and applies the param mapper's
/// `on_load` transformation.
    pub fn transform_for_load(self, tensor: Tensor<B, D, Int>, param_id: ParamId) -> Self {
// Make sure we load the tensor into the same module device.
/// Transform a parameter for saving by applying save transformations.
///
/// This method is used to prepare a parameter for saving (typically during serialization).
/// It applies the param mapper's `on_save` transformation, which can be used
/// to modify the tensor before serialization (e.g., quantization, precision conversion).
    pub fn transform_for_save(&self) -> Self {
impl<B: Backend, const D: usize> Param<Tensor<B, D, Bool>> {
/// The shape of the parameter, **without triggering initialization**.
///
/// This is critical for shape validation during loading: when applying tensors to an
/// uninitialized parameter, we need to validate the shape without triggering the
/// initialization function (which would allocate an unnecessary tensor).
///
/// **Returns:**
/// - For uninitialized params: the shape from the `Uninitialized` struct
/// - For initialized params: the actual shape from the tensor
///
/// Use this instead of [crate::tensor::Tensor::shape] when you need the shape but want to
/// preserve lazy initialization.
    pub fn lazy_shape(&self) -> burn_tensor::Shape {
/// Transform a parameter for loading by applying load transformations.
///
/// This method is used to restore a parameter from a tensor (typically during deserialization).
/// It ensures the tensor is moved to the expected device and applies the param mapper's
/// `on_load` transformation.
    pub fn transform_for_load(self, tensor: Tensor<B, D, Bool>, param_id: ParamId) -> Self {
// Make sure we load the tensor into the same module device.
/// Transform a parameter for saving by applying save transformations.
///
/// This method is used to prepare a parameter for saving (typically during serialization).
/// It applies the param mapper's `on_save` transformation, which can be used
/// to modify the tensor before serialization (e.g., quantization, precision conversion).
    pub fn transform_for_save(&self) -> Self {
impl<const D: usize, B: Backend> Module<B> for Param<Tensor<B, D>> {
    type Record = Param<Tensor<B, D>>;
    fn visit<V: ModuleVisitor<B>>(&self, visitor: &mut V) {
    fn map<M: ModuleMapper<B>>(self, mapper: &mut M) -> Self {
    fn into_record(self) -> Self::Record {
    fn load_record(self, record: Self::Record) -> Self {
    fn to_device(self, device: &Device<B>) -> Self {
    fn fork(self, device: &Device<B>) -> Self {
    fn collect_devices(&self, mut devices: Vec<Device<B>>) -> Vec<Device<B>> {
impl<const D: usize, B: Backend> ModuleDisplayDefault for Param<Tensor<B, D>> {
    fn content(&self, content: Content) -> Option<Content> {
impl<const D: usize, B: Backend> ModuleDisplay for Param<Tensor<B, D>> {
impl<const D: usize, B: Backend> Module<B> for Param<Tensor<B, D, Int>> {
    type Record = Param<Tensor<B, D, Int>>;
    fn visit<V: ModuleVisitor<B>>(&self, visitor: &mut V) {
    fn map<M: ModuleMapper<B>>(self, mapper: &mut M) -> Self {
    fn into_record(self) -> Self::Record {
    fn load_record(self, record: Self::Record) -> Self {
    fn to_device(self, device: &Device<B>) -> Self {
    fn fork(self, device: &Device<B>) -> Self {
// Don't support autodiff.
    fn collect_devices(&self, mut devices: Vec<Device<B>>) -> Vec<Device<B>> {
impl<const D: usize, B: Backend> ModuleDisplayDefault for Param<Tensor<B, D, Int>> {
    fn content(&self, content: Content) -> Option<Content> {
impl<const D: usize, B: Backend> ModuleDisplay for Param<Tensor<B, D, Int>> {
impl<const D: usize, B: Backend> Module<B> for Param<Tensor<B, D, Bool>> {
    type Record = Param<Tensor<B, D, Bool>>;
    fn visit<V: ModuleVisitor<B>>(&self, visitor: &mut V) {
    fn map<M: ModuleMapper<B>>(self, mapper: &mut M) -> Self {
    fn into_record(self) -> Self::Record {
    fn load_record(self, record: Self::Record) -> Self {
    fn to_device(self, device: &Device<B>) -> Self {
    fn fork(self, device: &Device<B>) -> Self {
// Don't support autodiff.
    fn collect_devices(&self, mut devices: Vec<Device<B>>) -> Vec<Device<B>> {
impl<const D: usize, B: Backend> ModuleDisplayDefault for Param<Tensor<B, D, Bool>> {
    fn content(&self, content: Content) -> Option<Content> {
impl<const D: usize, B: Backend> ModuleDisplay for Param<Tensor<B, D, Bool>> {
impl<const D: usize, B: AutodiffBackend> AutodiffModule<B> for Param<Tensor<B, D>> {
    type InnerModule = Param<Tensor<B::InnerBackend, D>>;
    fn valid(&self) -> Self::InnerModule {
impl<const D: usize, B: AutodiffBackend> AutodiffModule<B> for Param<Tensor<B, D, Int>> {
    type InnerModule = Param<Tensor<B::InnerBackend, D, Int>>;
    fn valid(&self) -> Self::InnerModule {
impl<const D: usize, B: AutodiffBackend> AutodiffModule<B> for Param<Tensor<B, D, Bool>> {
    type InnerModule = Param<Tensor<B::InnerBackend, D, Bool>>;
    fn valid(&self) -> Self::InnerModule {
    fn test_load_record_setting() {
struct ParamIdCollector<'a, M> {
impl<B, M> ModuleVisitor<B> for ParamIdCollector<'_, M>
where
    B: Backend,
    M: Module<B>,
{
    fn visit_float<const D: usize>(&mut self, param: &Param<Tensor<B, D>>) {
    fn visit_int<const D: usize>(&mut self, param: &Param<Tensor<B, D, Int>>) {
    fn visit_bool<const D: usize>(&mut self, param: &Param<Tensor<B, D, Bool>>) {
/// List all the parameter ids in a module.
pub fn list_param_ids<M: Module<B>, B: Backend>(module: &M) -> Vec<ParamId> {
/// Describes how to quantize a module.
pub struct Quantizer {
/// The calibration method used in quantization.
/// The quantization scheme.
impl<B: Backend> ModuleMapper<B> for Quantizer {
    fn map_float<const D: usize>(&mut self, param: Param<Tensor<B, D>>) -> Param<Tensor<B, D>> {
    type B = TestBackend;
    fn should_quantize_module() {
/// Overrides float and int tensors of [burn modules](super::Module).
///
/// This is useful for testing.
pub struct Reinitializer<B: Backend> {
enum ReinitStrategy<E> {
impl<B: Backend> Default for Reinitializer<B> {
    fn default() -> Self {
impl<B: Backend> Reinitializer<B> {
/// Create a new [reinitializer](Reinitializer).
    pub fn new() -> Self {
/// Apply the reinitialization to the given [module](Module).
    pub fn apply<M: Module<B>>(mut self, module: M) -> M {
/// Set the reinitialization strategy to constant for all tensors.
    pub fn constant(self, constant: f64) -> Self {
/// Set the reinitialization strategy to constant for float tensors.
    pub fn constant_float(mut self, constant: f64) -> Self {
/// Set the reinitialization strategy to constant for int tensors.
    pub fn constant_int(mut self, constant: i64) -> Self {
/// Set the reinitialization strategy to random for all tensors.
    pub fn random(self, seed: u64, min: f64, max: f64) -> Self {
/// Set the reinitialization strategy to random for float tensors.
    pub fn random_float(mut self, seed: u64, min: f64, max: f64) -> Self {
/// Set the reinitialization strategy to random for int tensors.
    pub fn random_int(mut self, seed: u64, min: i64, max: i64) -> Self {
/// Set the reinitialization strategy to range for all tensors.
    pub fn range(self, min: f64, max: f64) -> Self {
/// Set the reinitialization strategy to range for float tensors.
    pub fn range_float(mut self, min: f64, max: f64) -> Self {
/// Set the reinitialization strategy to range for int tensors.
    pub fn range_int(mut self, min: i64, max: i64) -> Self {
impl<B: Backend> ModuleMapper<B> for Reinitializer<B> {
    fn map_float<const D: usize>(
        &mut self,
        param: super::Param<Tensor<B, D>>,
    ) -> super::Param<Tensor<B, D>> {
    fn map_int<const D: usize>(
        &mut self,
        param: super::Param<Tensor<B, D, burn_tensor::Int>>,
    ) -> super::Param<Tensor<B, D, burn_tensor::Int>> {
    fn map_bool<const D: usize>(
        &mut self,
        param: super::Param<Tensor<B, D, burn_tensor::Bool>>,
    ) -> super::Param<Tensor<B, D, burn_tensor::Bool>> {
fn resolve<E: Element>(min: E, max: E, num_elements: usize) -> (E, E) {
fn random_vector<E: Element>(seed: u64, min: f64, max: f64, num_elements: usize) -> Vec<E> {
/// Trait to define a family of types which can be recorded using any [settings](PrecisionSettings).
pub trait Record<B: Backend>: Send {
/// Type of the item that can be serialized and deserialized.
    type Item<S: PrecisionSettings>: Serialize + DeserializeOwned;
/// Convert the current record into the corresponding item that follows the given [settings](PrecisionSettings).
    fn into_item<S: PrecisionSettings>(self) -> Self::Item<S>;
/// Convert the given item into a record.
    fn from_item<S: PrecisionSettings>(item: Self::Item<S>, device: &B::Device) -> Self;
/// Recorder trait specialized to save and load data to and from files.
pub trait FileRecorder<B: Backend>:
    Recorder<B, RecordArgs = PathBuf, RecordOutput = (), LoadArgs = PathBuf>
{
/// File extension of the format used by the recorder.
    fn file_extension() -> &'static str;
/// Default [file recorder](FileRecorder).
pub type DefaultFileRecorder<S> = NamedMpkFileRecorder<S>;
/// File recorder using the [bincode format](bincode).
pub struct BinFileRecorder<S: PrecisionSettings> {
/// File recorder using the [bincode format](bincode) compressed with gzip.
pub struct BinGzFileRecorder<S: PrecisionSettings> {
/// File recorder using the [json format](serde_json) compressed with gzip.
pub struct JsonGzFileRecorder<S: PrecisionSettings> {
/// File recorder using [pretty json format](serde_json) for easy readability.
pub struct PrettyJsonFileRecorder<S: PrecisionSettings> {
/// File recorder using the [named msgpack](rmp_serde) format compressed with gzip.
pub struct NamedMpkGzFileRecorder<S: PrecisionSettings> {
/// File recorder using the [named msgpack](rmp_serde) format.
pub struct NamedMpkFileRecorder<S: PrecisionSettings> {
impl<S: PrecisionSettings, B: Backend> FileRecorder<B> for BinGzFileRecorder<S> {
    fn file_extension() -> &'static str {
impl<S: PrecisionSettings, B: Backend> FileRecorder<B> for BinFileRecorder<S> {
    fn file_extension() -> &'static str {
impl<S: PrecisionSettings, B: Backend> FileRecorder<B> for JsonGzFileRecorder<S> {
    fn file_extension() -> &'static str {
impl<S: PrecisionSettings, B: Backend> FileRecorder<B> for PrettyJsonFileRecorder<S> {
    fn file_extension() -> &'static str {
impl<S: PrecisionSettings, B: Backend> FileRecorder<B> for NamedMpkGzFileRecorder<S> {
    fn file_extension() -> &'static str {
impl<S: PrecisionSettings, B: Backend> FileRecorder<B> for NamedMpkFileRecorder<S> {
    fn file_extension() -> &'static str {
// Add parent directories if they don't exist
impl<S: PrecisionSettings, B: Backend> Recorder<B> for BinGzFileRecorder<S> {
    type Settings = S;
    type RecordArgs = PathBuf;
    type RecordOutput = ();
    type LoadArgs = PathBuf;
    fn save_item<I: Serialize>(
        &self,
        item: I,
        mut file: Self::RecordArgs,
    ) -> Result<(), RecorderError> {
    fn load_item<I: DeserializeOwned>(
        &self,
        file: &mut Self::LoadArgs,
    ) -> Result<I, RecorderError> {
impl<S: PrecisionSettings, B: Backend> Recorder<B> for BinFileRecorder<S> {
    type Settings = S;
    type RecordArgs = PathBuf;
    type RecordOutput = ();
    type LoadArgs = PathBuf;
    fn save_item<I: Serialize>(
        &self,
        item: I,
        mut file: Self::RecordArgs,
    ) -> Result<(), RecorderError> {
    fn load_item<I: DeserializeOwned>(
        &self,
        file: &mut Self::LoadArgs,
    ) -> Result<I, RecorderError> {
impl<S: PrecisionSettings, B: Backend> Recorder<B> for JsonGzFileRecorder<S> {
    type Settings = S;
    type RecordArgs = PathBuf;
    type RecordOutput = ();
    type LoadArgs = PathBuf;
    fn save_item<I: Serialize>(
        &self,
        item: I,
        mut file: Self::RecordArgs,
    ) -> Result<(), RecorderError> {
    fn load_item<I: DeserializeOwned>(
        &self,
        file: &mut Self::LoadArgs,
    ) -> Result<I, RecorderError> {
impl<S: PrecisionSettings, B: Backend> Recorder<B> for PrettyJsonFileRecorder<S> {
    type Settings = S;
    type RecordArgs = PathBuf;
    type RecordOutput = ();
    type LoadArgs = PathBuf;
    fn save_item<I: Serialize>(
        &self,
        item: I,
        mut file: Self::RecordArgs,
    ) -> Result<(), RecorderError> {
    fn load_item<I: DeserializeOwned>(
        &self,
        file: &mut Self::LoadArgs,
    ) -> Result<I, RecorderError> {
impl<S: PrecisionSettings, B: Backend> Recorder<B> for NamedMpkGzFileRecorder<S> {
    type Settings = S;
    type RecordArgs = PathBuf;
    type RecordOutput = ();
    type LoadArgs = PathBuf;
    fn save_item<I: Serialize>(
        &self,
        item: I,
        mut file: Self::RecordArgs,
    ) -> Result<(), RecorderError> {
    fn load_item<I: DeserializeOwned>(
        &self,
        file: &mut Self::LoadArgs,
    ) -> Result<I, RecorderError> {
impl<S: PrecisionSettings, B: Backend> Recorder<B> for NamedMpkFileRecorder<S> {
    type Settings = S;
    type RecordArgs = PathBuf;
    type RecordOutput = ();
    type LoadArgs = PathBuf;
    fn save_item<I: Serialize>(
        &self,
        item: I,
        mut file: Self::RecordArgs,
    ) -> Result<(), RecorderError> {
    fn load_item<I: DeserializeOwned>(
        &self,
        file: &mut Self::LoadArgs,
    ) -> Result<I, RecorderError> {
    fn file_path() -> PathBuf {
    fn test_can_save_and_load_jsongz_format() {
    fn test_can_save_and_load_bin_format() {
    fn test_can_save_and_load_bingz_format() {
    fn test_can_save_and_load_pretty_json_format() {
    fn test_can_save_and_load_mpkgz_format() {
    fn test_can_save_and_load_mpk_format() {
    fn test_can_save_and_load<Recorder>(recorder: Recorder)
    where
        Recorder: FileRecorder<TestBackend>,
    {
    pub enum PaddingConfig2d {
// Dummy model with different record types
    pub struct Model<B: Backend> {
    pub fn create_model(device: &<TestBackend as Backend>::Device) -> Model<TestBackend> {
/// Recorder trait specialized to save and load data to and from bytes.
///
/// # Notes
///
/// This is especially useful in no_std environment where weights are stored directly in
/// compiled binaries.
pub trait BytesRecorder<
    B: Backend,
    L: AsRef<[u8]> + Send + Sync + core::fmt::Debug + Clone + core::default::Default,
>: Recorder<B, RecordArgs = (), RecordOutput = Vec<u8>, LoadArgs = L>
{
/// In memory recorder using the [bincode format](bincode).
pub struct BinBytesRecorder<
    S: PrecisionSettings,
    L: AsRef<[u8]> + Send + Sync + core::fmt::Debug + Clone + core::default::Default = Vec<u8>,
> {
impl<
    S: PrecisionSettings,
    B: Backend,
    L: AsRef<[u8]> + Send + Sync + core::fmt::Debug + Clone + core::default::Default,
> BytesRecorder<B, L> for BinBytesRecorder<S, L>
{
impl<
    S: PrecisionSettings,
    B: Backend,
    L: AsRef<[u8]> + Send + Sync + core::fmt::Debug + Clone + core::default::Default,
> Recorder<B> for BinBytesRecorder<S, L>
{
    type Settings = S;
    type RecordArgs = ();
    type RecordOutput = Vec<u8>;
    type LoadArgs = L;
    fn save_item<I: Serialize>(
        &self,
        item: I,
        _args: Self::RecordArgs,
    ) -> Result<Self::RecordOutput, RecorderError> {
    fn load_item<I: DeserializeOwned>(
        &self,
        args: &mut Self::LoadArgs,
    ) -> Result<I, RecorderError> {
/// In memory recorder using the [Named MessagePack](rmp_serde).
pub struct NamedMpkBytesRecorder<S: PrecisionSettings> {
impl<S: PrecisionSettings, B: Backend> BytesRecorder<B, Vec<u8>> for NamedMpkBytesRecorder<S> {
impl<S: PrecisionSettings, B: Backend> Recorder<B> for NamedMpkBytesRecorder<S> {
    type Settings = S;
    type RecordArgs = ();
    type RecordOutput = Vec<u8>;
    type LoadArgs = Vec<u8>;
    fn save_item<I: Serialize>(
        &self,
        item: I,
        _args: Self::RecordArgs,
    ) -> Result<Self::RecordOutput, RecorderError> {
    fn load_item<I: DeserializeOwned>(
        &self,
        args: &mut Self::LoadArgs,
    ) -> Result<I, RecorderError> {
    fn test_can_save_and_load_bin_format() {
    fn test_can_save_and_load_named_mpk_format() {
    fn test_can_save_and_load<Recorder>(recorder: Recorder)
    where
        Recorder: BytesRecorder<TestBackend, Vec<u8>>,
    {
    pub fn create_model<B: Backend>(device: &B::Device) -> SimpleLinear<B> {
impl<B> Record<B> for ()
where
    B: Backend,
{
    type Item<S: PrecisionSettings> = ();
    fn into_item<S: PrecisionSettings>(self) -> Self::Item<S> {
    fn from_item<S: PrecisionSettings>(_item: Self::Item<S>, _device: &B::Device) -> Self {
impl<T, B> Record<B> for Vec<T>
where
    T: Record<B>,
    B: Backend,
{
    type Item<S: PrecisionSettings> = Vec<T::Item<S>>;
    fn into_item<S: PrecisionSettings>(self) -> Self::Item<S> {
    fn from_item<S: PrecisionSettings>(item: Self::Item<S>, device: &B::Device) -> Self {
impl<T, B> Record<B> for Option<T>
where
    T: Record<B>,
    B: Backend,
{
    type Item<S: PrecisionSettings> = Option<T::Item<S>>;
    fn into_item<S: PrecisionSettings>(self) -> Self::Item<S> {
    fn from_item<S: PrecisionSettings>(item: Self::Item<S>, device: &B::Device) -> Self {
impl<const N: usize, T, B> Record<B> for [T;
/// The record item is an array of the record item of the elements.
/// The reason why we wrap the array in a struct is because serde does not support
/// deserializing arrays of variable size,
/// see [serde/issues/1937](https://github.com/serde-rs/serde/issues/1937).
/// for backward compatibility reasons. Serde APIs were created before const generics.
    type Item<S: PrecisionSettings> = Array<N, T::Item<S>>;
    fn into_item<S: PrecisionSettings>(self) -> Self::Item<S> {
    fn from_item<S: PrecisionSettings>(item: Self::Item<S>, device: &B::Device) -> Self {
/// A macro for generating implementations for tuple records of different sizes.
/// For example: `impl_record_tuple!([R0, R1][0, 1])`.
/// Would generate an implementation for a tuple of size 2.
/// For this macro to work properly, please adhere to the convention:
/// `impl_record_tuple!([R0, R1, ..., Rn][0, 1, ..., n])`.
// `$r` represents the generic records.
// `$i` represents the indices of the records in the tuple.
        impl<B, $($r,)*> Record<B> for ($($r,)*)
        where
            B: Backend,
            $($r: Record<B>),*
        {
            type Item<S: PrecisionSettings> = ($($r::Item<S>,)*);
            fn into_item<S: PrecisionSettings>(self) -> Self::Item<S> {
            fn from_item<S: PrecisionSettings>(item: Self::Item<S>, device: &B::Device) -> Self {
impl<T, B> Record<B> for HashMap<ParamId, T>
where
    T: Record<B>,
    B: Backend,
{
    type Item<S: PrecisionSettings> = HashMap<String, T::Item<S>>;
    fn into_item<S: PrecisionSettings>(self) -> Self::Item<S> {
    fn from_item<S: PrecisionSettings>(item: Self::Item<S>, device: &B::Device) -> Self {
/// (De)serialize parameters into a clean format.
pub struct ParamSerde<T> {
impl<B, const D: usize> Record<B> for Param<Tensor<B, D>>
where
    B: Backend,
{
    type Item<S: PrecisionSettings> = ParamSerde<FloatTensorSerde<S>>;
    fn into_item<S: PrecisionSettings>(self) -> Self::Item<S> {
    fn from_item<S: PrecisionSettings>(item: Self::Item<S>, device: &B::Device) -> Self {
// Same behavior as when we create a new
// Param from a tensor.
impl<B, const D: usize> Record<B> for Param<Tensor<B, D, Int>>
where
    B: Backend,
{
    type Item<S: PrecisionSettings> = ParamSerde<IntTensorSerde<S>>;
    fn into_item<S: PrecisionSettings>(self) -> Self::Item<S> {
    fn from_item<S: PrecisionSettings>(item: Self::Item<S>, device: &B::Device) -> Self {
impl<B, const D: usize> Record<B> for Param<Tensor<B, D, Bool>>
where
    B: Backend,
{
    type Item<S: PrecisionSettings> = ParamSerde<BoolTensorSerde>;
    fn into_item<S: PrecisionSettings>(self) -> Self::Item<S> {
    fn from_item<S: PrecisionSettings>(item: Self::Item<S>, device: &B::Device) -> Self {
// Type that can be serialized as is without any conversion.
    ($type:ty) => {
        impl<B: Backend> Record<B> for $type {
            type Item<S: PrecisionSettings> = $type;
            fn into_item<S: PrecisionSettings>(self) -> Self::Item<S> {
            fn from_item<S: PrecisionSettings>(item: Self::Item<S>, _device: &B::Device) -> Self {
// General Types
// Float Types
// Unsigned Integer Types
// Signed Integer Types
/// A wrapper around an array of size N, so that it can be serialized and deserialized
/// using serde.
///
/// The reason why we wrap the array in a struct is because serde does not support
/// deserializing arrays of variable size,
/// see [serde/issues/1937](https://github.com/serde-rs/serde/issues/1937)
/// for backward compatibility reasons. Serde APIs were created before const generics.
pub struct Array<const N: usize, T>([T;
impl<T: Serialize, const N: usize> Serialize for Array<N, T> {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: serde::Serializer,
    {
impl<'de, T, const N: usize> Deserialize<'de> for Array<N, T>
where
    T: Deserialize<'de>,
{
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        struct ArrayVisitor<T, const N: usize> {
        impl<'de, T, const N: usize> Visitor<'de> for ArrayVisitor<T, N>
        where
            T: Deserialize<'de>,
        {
            type Value = Array<N, T>;
            fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
            fn visit_seq<A>(self, mut seq: A) -> Result<Self::Value, A::Error>
            where
                A: SeqAccess<'de>,
            {
/// Record any item implementing [Serialize](Serialize) and [DeserializeOwned](DeserializeOwned).
pub trait Recorder<B: Backend>:
    Send + Sync + core::default::Default + core::fmt::Debug + Clone
{
/// Type of the settings used by the recorder.
    type Settings: PrecisionSettings;
/// Arguments used to record objects.
    type RecordArgs: Clone;
/// Record output type.
    type RecordOutput;
/// Arguments used to load recorded objects.
    type LoadArgs;
/// Records an item.
///
/// # Arguments
///
/// * `record` - The item to record.
/// * `args` - Arguments used to record the item.
///
/// # Returns
///
/// The output of the recording.
    fn record<R>(
        &self,
        record: R,
        args: Self::RecordArgs,
    ) -> Result<Self::RecordOutput, RecorderError>
    where
        R: Record<B>,
    {
/// Load an item from the given arguments.
    fn load<R>(&self, mut args: Self::LoadArgs, device: &B::Device) -> Result<R, RecorderError>
    where
        R: Record<B>,
    {
/// Saves an item.
///
/// This method is used by [record](Recorder::record) to save the item.
///
/// # Arguments
///
/// * `item` - Item to save.
/// * `args` - Arguments to use to save the item.
///
/// # Returns
///
/// The output of the save operation.
    fn save_item<I: Serialize>(
        &self,
        item: I,
        args: Self::RecordArgs,
    ) -> Result<Self::RecordOutput, RecorderError>;
/// Loads an item.
///
/// This method is used by [load](Recorder::load) to load the item.
///
/// # Arguments
///
/// * `args` - Arguments to use to load the item.
///
/// # Returns
///
/// The loaded item.
    fn load_item<I>(&self, args: &mut Self::LoadArgs) -> Result<I, RecorderError>
    where
        I: DeserializeOwned;
fn recorder_metadata<R, B>() -> BurnMetadata
where
    R: Recorder<B>,
    B: Backend,
{
/// Error that can occur when using a [Recorder](Recorder).
pub enum RecorderError {
/// File not found.
/// Failed to read file.
/// Other error.
impl core::fmt::Display for RecorderError {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
impl core::error::Error for RecorderError {
pub(crate) fn bin_config() -> bincode::config::Configuration {
/// Metadata of a record.
pub struct BurnMetadata {
/// Float type used to record the item.
/// Int type used to record the item.
/// Format used to record the item.
/// Burn record version used to record the item.
/// Settings used to record the item.
/// Record that can be saved by a [Recorder](Recorder).
pub struct BurnRecord<I, B: Backend> {
/// Metadata of the record.
/// Item to record.
impl<I, B: Backend> BurnRecord<I, B> {
/// Creates a new record.
///
/// # Arguments
///
/// * `item` - Item to record.
///
/// # Returns
///
/// The new record.
    pub fn new<R: Recorder<B>>(item: I) -> Self {
/// Record that can be saved by a [Recorder](Recorder) without the item.
pub struct BurnRecordNoItem {
/// Metadata of the record.
/// Default recorder.
///
/// It uses the [named msgpack](rmp_serde) format for serialization with full precision.
pub type DefaultRecorder = DefaultFileRecorder<FullPrecisionSettings>;
/// Recorder optimized for compactness.
///
/// It uses the [named msgpack](rmp_serde) format for serialization with half precision.
/// If you are looking for the recorder that offers the smallest file size, have a look at
/// [sensitive compact recorder](SensitiveCompactRecorder).
pub type CompactRecorder = DefaultFileRecorder<HalfPrecisionSettings>;
/// Recorder optimized for compactness making it a good choice for model deployment.
///
/// It uses the [bincode](bincode) format for serialization and half precision.
/// This format is not resilient to type changes since no metadata is encoded.
/// Favor [default recorder](DefaultRecorder) or [compact recorder](CompactRecorder)
/// for long term data storage.
pub type SensitiveCompactRecorder = BinGzFileRecorder<HalfPrecisionSettings>;
/// Training recorder compatible with no-std inference.
pub type NoStdTrainingRecorder = BinFileRecorder<FullPrecisionSettings>;
/// Inference recorder compatible with no-std.
pub type NoStdInferenceRecorder = BinBytesRecorder<FullPrecisionSettings, &'static [u8]>;
/// Debug recorder.
///
/// It uses the [pretty json](serde_json) format for serialization with full precision making it
/// human readable.
pub type DebugRecordSettings = PrettyJsonFileRecorder<FullPrecisionSettings>;
    static FILE_PATH: &str = "/tmp/burn_test_record";
    fn err_when_invalid_item() {
        struct Item<S: PrecisionSettings> {
        impl<D, B> Record<B> for Item<D>
        where
            D: PrecisionSettings,
            B: Backend,
        {
            type Item<S: PrecisionSettings> = Item<S>;
            fn into_item<S: PrecisionSettings>(self) -> Self::Item<S> {
            fn from_item<S: PrecisionSettings>(item: Self::Item<S>, _device: &B::Device) -> Self {
// Serialize in f32.
// Can't deserialize f32 into f16.
/// A trait that defines the adapter for a Burn module.
///
/// This is used to adapt an incoming module to a Burn module.
pub trait BurnModuleAdapter: Sized {
/// Adapts a module.
    fn adapt(name: &str, data: NestedValue) -> NestedValue {
/// Adapts a linear module.
    fn adapt_linear(data: NestedValue) -> NestedValue {
/// Adapts a Convolution 1D module.
    fn adapt_conv1d(data: NestedValue) -> NestedValue {
/// Adapts a Convolution 2D module.
    fn adapt_conv2d(data: NestedValue) -> NestedValue {
/// Adapts a Convolution 3D module.
    fn adapt_conv3d(data: NestedValue) -> NestedValue {
/// Adapts convolution transpose 1D module.
    fn adapt_conv_transpose_1d(data: NestedValue) -> NestedValue {
/// Adapts convolution transpose 2D module.
    fn adapt_conv_transpose_2d(data: NestedValue) -> NestedValue {
/// Adapts convolution transpose 2D module.
    fn adapt_conv_transpose_3d(data: NestedValue) -> NestedValue {
/// Adapts embedding module.
    fn adapt_embedding(data: NestedValue) -> NestedValue {
/// Adapts group normalization module.
    fn adapt_group_norm(data: NestedValue) -> NestedValue {
/// Adapts layer normalization module.
    fn adapt_layer_norm(data: NestedValue) -> NestedValue {
/// Adapts batch normalization module.
    fn adapt_batch_norm(data: NestedValue) -> NestedValue {
/// Default adapter that takes no action.
pub struct DefaultAdapter;
impl BurnModuleAdapter for DefaultAdapter {
/// The main data structure used for deserialization.
///
/// It can hold tree-like structures of nested maps and vectors.
pub enum NestedValue {
/// The default value, which actually does not hold any value and it is used to indicate that
/// the value should be populated with the default value. It contains an optional string with
/// the originator field name.
/// A boolean value.
/// A string value.
/// Floating point 32-bit value.
/// Floating point 64-bit value.
/// Signed 16-bit integer value.
/// Signed 32-bit integer value.
/// Signed 64-bit integer value.
/// Unsigned 8-bit integer value.
/// Unsigned 16-bit integer value used for bf16 and f16 serialization
/// Unsigned 64-bit integer value.
/// A map of nested values (typically used for structs)
/// A vector of nested values (typically used for vector of structs or numbers)
/// A vector of 8-bit unsigned integer values.
/// A vector of 16-bit unsigned integer values.
/// A vector of 32-bit floating point values.
/// An opaque vector of bytes, with alignment.
impl NestedValue {
/// Get the nested value as a map.
    pub fn as_map(self) -> Option<HashMap<String, NestedValue>> {
/// Get the nested value as a boolean.
    pub fn as_bool(self) -> Option<bool> {
/// Get the nested value as a string.
    pub fn as_string(self) -> Option<String> {
/// Get the nested value as a f32.
    pub fn as_f32(self) -> Option<f32> {
/// Get the nested value as a f64.
    pub fn as_f64(self) -> Option<f64> {
/// Get the nested value as an i16.
    pub fn as_i16(self) -> Option<i16> {
/// Get the nested value as an i32.
    pub fn as_i32(self) -> Option<i32> {
/// Get the nested value as an i64.
    pub fn as_i64(self) -> Option<i64> {
/// Get the nested value as a u8.
    pub fn as_u8(self) -> Option<u8> {
/// Get the nested value as a u16.
    pub fn as_u16(self) -> Option<u16> {
/// Get the nested value as a u64.
    pub fn as_u64(self) -> Option<u64> {
/// Get the nested value as a vector of bytes.
    pub fn as_bytes(self) -> Option<Bytes> {
/// Deserialize a nested value into a record type.
    pub fn try_into_record<T, PS, A, B>(self, device: &B::Device) -> Result<T, Error>
    where
        B: Backend,
        T: Record<B>,
        PS: PrecisionSettings,
        A: BurnModuleAdapter,
    {
// Convert the deserialized item into a Record instance
/// Remap the tensor locations according to the key remapping.
///
/// # Arguments
///
/// * `tensors` - A map of tensors.
/// * `key_remap` - A vector of tuples containing a regular expression and a replacement string.
///   See [regex::Regex::replace](https://docs.rs/regex/latest/regex/struct.Regex.html#method.replace)
///   for more information.
///
/// # Returns
///
/// A map of tensors with the remapped keys and
/// a vector of tuples containing the remapped and original.
pub fn remap<T>(
    mut tensors: HashMap<String, T>,
    key_remap: Vec<(Regex, String)>,
) -> (HashMap<String, T>, Vec<(String, String)>) {
// Name is the same as the remapped name
/// Helper function to insert a value into a nested map/vector of tensors.
fn insert_nested_value(current: &mut NestedValue, keys: &[&str], value: NestedValue) {
/// A trait for encapsulating the serialization logic.
pub trait Serializable {
/// Serializes the object into a `NestedValue` using the provided `Serializer`.
/// This method is generic over the precision settings `PS`.
///
/// # Parameters
/// - `serializer`: The `Serializer` to use for serializing the object.
///
/// # Returns
/// - `Result<NestedValue, Error>`: The result of serialization.
///   Returns a `NestedValue` on success,
///   or an `Error` on failure.
///
/// # Type Parameters
/// - `PS`: The precision settings to use during serialization.
///   This is a generic parameter and can be any type
///   that implements the `PrecisionSettings` trait.
    fn serialize<PS>(&self, serializer: Serializer) -> Result<NestedValue, Error>
    where
        PS: PrecisionSettings;
/// Convert a vector of tensors to a nested value.
pub fn unflatten<PS, T>(input: HashMap<String, T>) -> Result<NestedValue, Error>
where
    PS: PrecisionSettings,
    T: Serializable,
{
/// Removes empty maps from the nested value.
///
/// We need to clean up empty maps from the nested value
/// in some cases when there is non-contiguous indices in keys.
fn cleanup_empty_maps(current: &mut NestedValue) {
fn write_vec_truncated<T: core::fmt::Debug>(
    vec: &[T],
    f: &mut core::fmt::Formatter,
) -> fmt::Result {
impl fmt::Debug for NestedValue {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
// Truncate values for vector
// Handle other variants as usual
const RECORD_ITEM_SUFFIX: &str = "RecordItem";
/// A deserializer for the nested value data structure.
pub struct Deserializer<A: BurnModuleAdapter> {
// This string starts with the input data and characters are truncated off
// the beginning as data is parsed.
impl<A: BurnModuleAdapter> Deserializer<A> {
/// Creates a new deserializer with the given nested value.
///
/// # Arguments
///
/// * `value` - A nested value.
/// * `default_for_missing_fields` - A boolean indicating whether to add missing fields with default value.
    pub fn new(value: NestedValue, default_for_missing_fields: bool) -> Self {
impl<'de, A: BurnModuleAdapter> serde::Deserializer<'de> for Deserializer<A> {
    type Error = Error;
    fn deserialize_any<V>(self, _visitor: V) -> Result<V::Value, Self::Error>
    where
        V: Visitor<'de>,
    {
    fn deserialize_struct<V>(
        self,
        name: &'static str,
        fields: &'static [&'static str],
        visitor: V,
    ) -> Result<V::Value, Self::Error>
    where
        V: Visitor<'de>,
    {
// Adapt modules
// Add missing fields into the map with default value if needed.
    fn deserialize_string<V>(self, visitor: V) -> Result<V::Value, Self::Error>
    where
        V: Visitor<'de>,
    {
    fn deserialize_ignored_any<V>(self, visitor: V) -> Result<V::Value, Self::Error>
    where
        V: Visitor<'de>,
    {
    fn deserialize_map<V>(self, visitor: V) -> Result<V::Value, Self::Error>
    where
        V: Visitor<'de>,
    {
    fn deserialize_bool<V>(self, visitor: V) -> Result<V::Value, Self::Error>
    where
        V: Visitor<'de>,
    {
    fn deserialize_i8<V>(self, _visitor: V) -> Result<V::Value, Self::Error>
    where
        V: Visitor<'de>,
    {
    fn deserialize_i16<V>(self, visitor: V) -> Result<V::Value, Self::Error>
    where
        V: Visitor<'de>,
    {
    fn deserialize_i32<V>(self, visitor: V) -> Result<V::Value, Self::Error>
    where
        V: Visitor<'de>,
    {
    fn deserialize_i64<V>(self, visitor: V) -> Result<V::Value, Self::Error>
    where
        V: Visitor<'de>,
    {
    fn deserialize_u8<V>(self, visitor: V) -> Result<V::Value, Self::Error>
    where
        V: Visitor<'de>,
    {
    fn deserialize_u16<V>(self, visitor: V) -> Result<V::Value, Self::Error>
    where
        V: Visitor<'de>,
    {
    fn deserialize_u32<V>(self, _visitor: V) -> Result<V::Value, Self::Error>
    where
        V: Visitor<'de>,
    {
    fn deserialize_u64<V>(self, visitor: V) -> Result<V::Value, Self::Error>
    where
        V: Visitor<'de>,
    {
    fn deserialize_f32<V>(self, visitor: V) -> Result<V::Value, Self::Error>
    where
        V: Visitor<'de>,
    {
    fn deserialize_f64<V>(self, visitor: V) -> Result<V::Value, Self::Error>
    where
        V: Visitor<'de>,
    {
    fn deserialize_char<V>(self, _visitor: V) -> Result<V::Value, Self::Error>
    where
        V: Visitor<'de>,
    {
    fn deserialize_str<V>(self, visitor: V) -> Result<V::Value, Self::Error>
    where
        V: Visitor<'de>,
    {
    fn deserialize_bytes<V>(self, _visitor: V) -> Result<V::Value, Self::Error>
    where
        V: Visitor<'de>,
    {
    fn deserialize_byte_buf<V>(self, visitor: V) -> Result<V::Value, Self::Error>
    where
        V: Visitor<'de>,
    {
    fn deserialize_option<V>(self, visitor: V) -> Result<V::Value, Self::Error>
    where
        V: Visitor<'de>,
    {
    fn deserialize_unit<V>(self, _visitor: V) -> Result<V::Value, Self::Error>
    where
        V: Visitor<'de>,
    {
    fn deserialize_unit_struct<V>(
        self,
        _name: &'static str,
        _visitor: V,
    ) -> Result<V::Value, Self::Error>
    where
        V: Visitor<'de>,
    {
    fn deserialize_newtype_struct<V>(
        self,
        _name: &'static str,
        visitor: V,
    ) -> Result<V::Value, Self::Error>
    where
        V: Visitor<'de>,
    {
    fn deserialize_seq<V>(self, visitor: V) -> Result<V::Value, Self::Error>
    where
        V: Visitor<'de>,
    {
    fn deserialize_tuple<V>(self, _len: usize, _visitor: V) -> Result<V::Value, Self::Error>
    where
        V: Visitor<'de>,
    {
    fn deserialize_tuple_struct<V>(
        self,
        _name: &'static str,
        _len: usize,
        _visitor: V,
    ) -> Result<V::Value, Self::Error>
    where
        V: Visitor<'de>,
    {
/// Deserializes an enum by attempting to match its variants against the provided data.
///
/// This function attempts to deserialize an enum by iterating over its possible variants
/// and trying to deserialize the data into each until one succeeds. We need to do this
/// because we don't have a way to know which variant to deserialize from the data.
///
/// This is similar to Serde's
/// [untagged enum deserialization](https://serde.rs/enum-representations.html#untagged),
/// but it's on the deserializer side. Using `#[serde(untagged)]` on the enum will force
/// using `deserialize_any`, which is not what we want because we want to use methods, such
/// as `visit_struct`. Also we do not wish to use auto generate code for Deserialize just
/// for enums because it will affect other serialization and deserialization, such
/// as JSON and Bincode.
///
/// # Safety
/// The function uses an unsafe block to clone the `visitor`. This is necessary because
/// the `Visitor` trait does not have a `Clone` implementation, and we need to clone it
/// as we are going to use it multiple times. The Visitor is a code generated unit struct
/// with no states or mutations, so it is safe to clone it in this case. We mainly care
/// about the `visit_enum` method, which is the only method that will be called on the
/// cloned visitor.
    fn deserialize_enum<V>(
        self,
        _name: &'static str,
        variants: &'static [&'static str],
        visitor: V,
    ) -> Result<V::Value, Self::Error>
    where
        V: Visitor<'de>,
    {
        fn clone_unsafely<T>(thing: &T) -> T {
// Allocate memory for the clone.
// Get a mutable pointer to the allocated memory.
// Copy the memory
                ptr::copy_nonoverlapping(thing as *const T, clone_ptr, 1);
// Assume the cloned data is initialized and convert it to an owned instance of T.
// Try each variant in order
// clone visitor to avoid moving it
    fn deserialize_identifier<V>(self, _visitor: V) -> Result<V::Value, Self::Error>
    where
        V: Visitor<'de>,
    {
/// A sequence access for a vector in the nested value data structure.
struct VecSeqAccess<A: BurnModuleAdapter, I> {
// Concrete implementation for `Vec<NestedValue>`
impl<A: BurnModuleAdapter> VecSeqAccess<A, NestedValue> {
    fn new(vec: NestedValue, default_for_missing_fields: bool) -> Self {
// Concrete implementation for `Vec<u8>`
impl<A: BurnModuleAdapter> VecSeqAccess<A, u8> {
    fn new(vec: NestedValue, default_for_missing_fields: bool) -> Self {
// Concrete implementation for `Vec<u16>`
impl<A: BurnModuleAdapter> VecSeqAccess<A, u16> {
    fn new(vec: NestedValue, default_for_missing_fields: bool) -> Self {
// Concrete implementation for `Vec<f32>`
impl<A: BurnModuleAdapter> VecSeqAccess<A, f32> {
    fn new(vec: NestedValue, default_for_missing_fields: bool) -> Self {
// Concrete implementation for `Vec<NestedValue>`
impl<'de, A> SeqAccess<'de> for VecSeqAccess<A, NestedValue>
where
    NestedValueWrapper<A>: IntoDeserializer<'de, Error>,
    A: BurnModuleAdapter,
{
    type Error = Error;
    fn next_element_seed<T>(&mut self, seed: T) -> Result<Option<T::Value>, Self::Error>
    where
        T: DeserializeSeed<'de>,
    {
// Concrete implementation for `Vec<u8>`
impl<'de, A> SeqAccess<'de> for VecSeqAccess<A, u8>
where
    NestedValueWrapper<A>: IntoDeserializer<'de, Error>,
    A: BurnModuleAdapter,
{
    type Error = Error;
    fn next_element_seed<T>(&mut self, seed: T) -> Result<Option<T::Value>, Self::Error>
    where
        T: DeserializeSeed<'de>,
    {
// Concrete implementation for `Vec<u16>`
impl<'de, A> SeqAccess<'de> for VecSeqAccess<A, u16>
where
    NestedValueWrapper<A>: IntoDeserializer<'de, Error>,
    A: BurnModuleAdapter,
{
    type Error = Error;
    fn next_element_seed<T>(&mut self, seed: T) -> Result<Option<T::Value>, Self::Error>
    where
        T: DeserializeSeed<'de>,
    {
// Concrete implementation for `Vec<f32>`
impl<'de, A> SeqAccess<'de> for VecSeqAccess<A, f32>
where
    NestedValueWrapper<A>: IntoDeserializer<'de, Error>,
    A: BurnModuleAdapter,
{
    type Error = Error;
    fn next_element_seed<T>(&mut self, seed: T) -> Result<Option<T::Value>, Self::Error>
    where
        T: DeserializeSeed<'de>,
    {
/// A map access for a map in the nested value data structure.
struct HashMapAccess<A: BurnModuleAdapter> {
impl<A: BurnModuleAdapter> HashMapAccess<A> {
    fn new(map: HashMap<String, NestedValue>, default_for_missing_fields: bool) -> Self {
impl<'de, A> MapAccess<'de> for HashMapAccess<A>
where
    String: IntoDeserializer<'de, Error>,
    NestedValueWrapper<A>: IntoDeserializer<'de, Error>,
    A: BurnModuleAdapter,
{
    type Error = Error;
    fn next_key_seed<T>(&mut self, seed: T) -> Result<Option<T::Value>, Self::Error>
    where
        T: DeserializeSeed<'de>,
    {
// Keep the value for the next call to next_value_seed.
// Deserialize the key.
    fn next_value_seed<T>(&mut self, seed: T) -> Result<T::Value, Self::Error>
    where
        T: DeserializeSeed<'de>,
    {
struct ProbeEnumAccess<A: BurnModuleAdapter> {
impl<A: BurnModuleAdapter> ProbeEnumAccess<A> {
    fn new(value: NestedValue, current_variant: String, default_for_missing_fields: bool) -> Self {
impl<'de, A> EnumAccess<'de> for ProbeEnumAccess<A>
where
    A: BurnModuleAdapter,
{
    type Error = Error;
    type Variant = Self;
    fn variant_seed<V>(self, seed: V) -> Result<(V::Value, Self::Variant), Self::Error>
    where
        V: DeserializeSeed<'de>,
    {
impl<'de, A> VariantAccess<'de> for ProbeEnumAccess<A>
where
    A: BurnModuleAdapter,
{
    type Error = Error;
    fn newtype_variant_seed<T>(self, seed: T) -> Result<T::Value, Self::Error>
    where
        T: DeserializeSeed<'de>,
    {
    fn unit_variant(self) -> Result<(), Self::Error> {
// Support tensor `DType` deserialization
// wrong match
    fn tuple_variant<V>(self, _len: usize, _visitor: V) -> Result<V::Value, Self::Error>
    where
        V: Visitor<'de>,
    {
    fn struct_variant<V>(
        self,
        _fields: &'static [&'static str],
        _visitor: V,
    ) -> Result<V::Value, Self::Error>
    where
        V: Visitor<'de>,
    {
/// A wrapper for the nested value data structure with a burn module adapter.
struct NestedValueWrapper<A: BurnModuleAdapter> {
impl<A: BurnModuleAdapter> NestedValueWrapper<A> {
    fn new(value: NestedValue, default_for_missing_fields: bool) -> Self {
impl<A: BurnModuleAdapter> IntoDeserializer<'_, Error> for NestedValueWrapper<A> {
    type Deserializer = Deserializer<A>;
    fn into_deserializer(self) -> Self::Deserializer {
/// A default deserializer that always returns the default value.
struct DefaultDeserializer {
/// The originator field name (the top-level missing field name)
impl DefaultDeserializer {
    fn new(originator_field_name: Option<String>) -> Self {
impl<'de> serde::Deserializer<'de> for DefaultDeserializer {
    type Error = Error;
    fn deserialize_any<V>(self, _visitor: V) -> Result<V::Value, Self::Error>
    where
        V: Visitor<'de>,
    {
    fn deserialize_i32<V>(self, visitor: V) -> Result<V::Value, Self::Error>
    where
        V: Visitor<'de>,
    {
    fn deserialize_f32<V>(self, visitor: V) -> Result<V::Value, Self::Error>
    where
        V: Visitor<'de>,
    {
    fn deserialize_i16<V>(self, visitor: V) -> Result<V::Value, Self::Error>
    where
        V: Visitor<'de>,
    {
    fn deserialize_i64<V>(self, visitor: V) -> Result<V::Value, Self::Error>
    where
        V: Visitor<'de>,
    {
    fn deserialize_u16<V>(self, visitor: V) -> Result<V::Value, Self::Error>
    where
        V: Visitor<'de>,
    {
    fn deserialize_u64<V>(self, visitor: V) -> Result<V::Value, Self::Error>
    where
        V: Visitor<'de>,
    {
    fn deserialize_f64<V>(self, visitor: V) -> Result<V::Value, Self::Error>
    where
        V: Visitor<'de>,
    {
    fn deserialize_bool<V>(self, visitor: V) -> Result<V::Value, Self::Error>
    where
        V: Visitor<'de>,
    {
    fn deserialize_char<V>(self, visitor: V) -> Result<V::Value, Self::Error>
    where
        V: Visitor<'de>,
    {
    fn deserialize_str<V>(self, visitor: V) -> Result<V::Value, Self::Error>
    where
        V: Visitor<'de>,
    {
    fn deserialize_i8<V>(self, visitor: V) -> Result<V::Value, Self::Error>
    where
        V: Visitor<'de>,
    {
    fn deserialize_u8<V>(self, visitor: V) -> Result<V::Value, Self::Error>
    where
        V: Visitor<'de>,
    {
    fn deserialize_u32<V>(self, visitor: V) -> Result<V::Value, Self::Error>
    where
        V: Visitor<'de>,
    {
    fn deserialize_option<V>(self, visitor: V) -> Result<V::Value, Self::Error>
    where
        V: Visitor<'de>,
    {
    fn deserialize_seq<V>(self, visitor: V) -> Result<V::Value, Self::Error>
    where
        V: Visitor<'de>,
    {
    fn deserialize_string<V>(self, visitor: V) -> Result<V::Value, Self::Error>
    where
        V: Visitor<'de>,
    {
    fn deserialize_struct<V>(
        self,
        name: &'static str,
        _fields: &'static [&'static str],
        _visitor: V,
    ) -> Result<V::Value, Self::Error>
    where
        V: Visitor<'de>,
    {
// Return an error if the originator field name is not set
    fn deserialize_tuple_struct<V>(
        self,
        _name: &'static str,
        len: usize,
        visitor: V,
    ) -> Result<V::Value, Self::Error>
    where
        V: Visitor<'de>,
    {
    fn deserialize_tuple<V>(self, len: usize, visitor: V) -> Result<V::Value, Self::Error>
    where
        V: Visitor<'de>,
    {
    fn deserialize_map<V>(self, visitor: V) -> Result<V::Value, Self::Error>
    where
        V: Visitor<'de>,
    {
        enum identifier ignored_any
    }
}

/// A default sequence access that always returns None (empty sequence).
pub struct DefaultSeqAccess {
impl Default for DefaultSeqAccess {
    fn default() -> Self {
impl DefaultSeqAccess {
/// Creates a new default sequence access with the given size hint.
    pub fn new(size: Option<usize>) -> Self {
impl<'de> SeqAccess<'de> for DefaultSeqAccess {
    type Error = Error;
    fn next_element_seed<T>(&mut self, seed: T) -> Result<Option<T::Value>, Self::Error>
    where
        T: DeserializeSeed<'de>,
    {
    fn size_hint(&self) -> Option<usize> {
/// A default map access that always returns None (empty map).
pub struct DefaultMapAccess;
impl Default for DefaultMapAccess {
    fn default() -> Self {
impl DefaultMapAccess {
/// Creates a new default map access.
    pub fn new() -> Self {
impl<'de> MapAccess<'de> for DefaultMapAccess {
    type Error = Error;
    fn next_key_seed<T>(&mut self, _seed: T) -> Result<Option<T::Value>, Self::Error>
    where
        T: DeserializeSeed<'de>,
    {
// Since this is a default implementation, we'll just return None.
    fn next_value_seed<T>(&mut self, _seed: T) -> Result<T::Value, Self::Error>
    where
        T: DeserializeSeed<'de>,
    {
    fn size_hint(&self) -> Option<usize> {
// Since this is a default implementation, we'll just return None.
/// The error type for Record serde.
pub enum Error {
/// Failed to deserialize.
/// Failed to serialize.
/// Encountered an invalid state.
/// Other error.
impl serde::de::Error for Error {
    fn custom<T: std::fmt::Display>(msg: T) -> Self {
impl serde::ser::Error for Error {
    fn custom<T: std::fmt::Display>(msg: T) -> Self {
// Implement From trait for Error to RecorderError
impl From<Error> for RecorderError {
    fn from(error: Error) -> Self {
//! Module contains the serde implementation for the record module
//! useful for custom importing model weights, such as PyTorch's pt file format.
/// The adapter trait that is used to convert the nested value to the module type.
/// The main data structure used for deserialization.
/// The deserializer that is used to convert the nested value to the record.
/// The deserializer that is used to convert the nested value to the record.
/// Error types.
/// Simple struct serializer that converts a struct into NestedValues.
///
/// NOTE: This is used to serialize Param structs into NestedValues and not so much for
/// the actual serialization of modules (although it could be used for that as well if all
/// primitive types are implemented).
pub struct Serializer {
/// The state of the serialization process
impl Serializer {
/// Creates a new serializer.
    pub fn new() -> Self {
impl Default for Serializer {
    fn default() -> Self {
impl SerializerTrait for Serializer {
    type Ok = NestedValue;
    type Error = Error;
    type SerializeSeq = Self;
    type SerializeTuple = ser::Impossible<NestedValue, Self::Error>;
    type SerializeTupleStruct = ser::Impossible<NestedValue, Self::Error>;
    type SerializeTupleVariant = ser::Impossible<NestedValue, Self::Error>;
    type SerializeMap = ser::Impossible<NestedValue, Self::Error>;
    type SerializeStruct = Self;
    type SerializeStructVariant = ser::Impossible<NestedValue, Self::Error>;
    fn serialize_struct(
        self,
        _name: &'static str,
        _len: usize,
    ) -> Result<Self::SerializeStruct, Self::Error> {
    fn serialize_newtype_struct<T>(
        self,
        _name: &'static str,
        value: &T,
    ) -> Result<Self::Ok, Self::Error>
    where
        T: Serialize + ?Sized,
    {
    fn serialize_seq(self, _len: Option<usize>) -> Result<Self::SerializeSeq, Self::Error> {
    fn serialize_i32(self, v: i32) -> Result<Self::Ok, Self::Error> {
    fn serialize_str(self, v: &str) -> Result<Self::Ok, Self::Error> {
    fn serialize_i16(self, v: i16) -> Result<Self::Ok, Self::Error> {
    fn serialize_i64(self, v: i64) -> Result<Self::Ok, Self::Error> {
    fn serialize_u16(self, v: u16) -> Result<Self::Ok, Self::Error> {
    fn serialize_u64(self, v: u64) -> Result<Self::Ok, Self::Error> {
    fn serialize_f32(self, v: f32) -> Result<Self::Ok, Self::Error> {
    fn serialize_f64(self, v: f64) -> Result<Self::Ok, Self::Error> {
// The following methods are not implemented because they are not needed for the
// serialization of Param structs.
    fn serialize_char(self, _v: char) -> Result<Self::Ok, Self::Error> {
    fn serialize_bytes(self, v: &[u8]) -> Result<Self::Ok, Self::Error> {
    fn serialize_none(self) -> Result<Self::Ok, Self::Error> {
    fn serialize_u32(self, _v: u32) -> Result<Self::Ok, Self::Error> {
    fn serialize_bool(self, _v: bool) -> Result<Self::Ok, Self::Error> {
    fn serialize_i8(self, _v: i8) -> Result<Self::Ok, Self::Error> {
    fn serialize_u8(self, v: u8) -> Result<Self::Ok, Self::Error> {
    fn serialize_some<T>(self, value: &T) -> Result<Self::Ok, Self::Error>
    where
        T: Serialize + ?Sized,
    {
    fn serialize_unit(self) -> Result<Self::Ok, Self::Error> {
    fn serialize_unit_struct(self, _name: &'static str) -> Result<Self::Ok, Self::Error> {
    fn serialize_unit_variant(
        self,
        _name: &'static str,
        _variant_index: u32,
        _variant: &'static str,
    ) -> Result<Self::Ok, Self::Error> {
    fn serialize_newtype_variant<T>(
        self,
        _name: &'static str,
        _variant_index: u32,
        _variant: &'static str,
        _value: &T,
    ) -> Result<Self::Ok, Self::Error>
    where
        T: Serialize + ?Sized,
    {
    fn serialize_tuple(self, _len: usize) -> Result<Self::SerializeTuple, Self::Error> {
    fn serialize_tuple_struct(
        self,
        _name: &'static str,
        _len: usize,
    ) -> Result<Self::SerializeTupleStruct, Self::Error> {
    fn serialize_tuple_variant(
        self,
        _name: &'static str,
        _variant_index: u32,
        _variant: &'static str,
        _len: usize,
    ) -> Result<Self::SerializeTupleVariant, Self::Error> {
    fn serialize_map(self, _len: Option<usize>) -> Result<Self::SerializeMap, Self::Error> {
    fn serialize_struct_variant(
        self,
        _name: &'static str,
        _variant_index: u32,
        _variant: &'static str,
        _len: usize,
    ) -> Result<Self::SerializeStructVariant, Self::Error> {
// Implementing the SerializeStruct trait for Serializer
impl SerializeStruct for Serializer {
    type Ok = NestedValue;
    type Error = Error;
    fn serialize_field<T>(&mut self, key: &'static str, value: &T) -> Result<(), Self::Error>
    where
        T: Serialize + ?Sized,
    {
// Inserting into the state
// Inserting into the state
    fn end(self) -> Result<Self::Ok, Self::Error> {
// If the state is empty, return an empty map
impl SerializeSeq for Serializer {
    type Ok = NestedValue;
    type Error = Error;
    fn serialize_element<T>(&mut self, value: &T) -> Result<(), Self::Error>
    where
        T: Serialize + ?Sized,
    {
// Inserting into the state
    fn end(self) -> Result<Self::Ok, Self::Error> {
// If the state is empty, return an empty vector
    struct MyStruct1 {
    struct MyStruct2 {
    struct MyStruct3 {
    fn test_serialize() {
// Compare the lengths of expected and actual serialized strings because
// the order of the fields is not guaranteed for HashMaps.
    fn test_param_serde() {
/// Settings allowing to control the precision when (de)serializing items.
pub trait PrecisionSettings:
    Send + Sync + core::fmt::Debug + core::default::Default + Clone
{
/// Float element type.
    type FloatElem: Element + Serialize + DeserializeOwned;
/// Integer element type.
    type IntElem: Element + Serialize + DeserializeOwned;
/// Default precision settings.
pub struct FullPrecisionSettings;
/// Precision settings optimized for compactness.
pub struct HalfPrecisionSettings;
/// Precision settings optimized for precision.
pub struct DoublePrecisionSettings;
impl PrecisionSettings for FullPrecisionSettings {
    type FloatElem = f32;
    type IntElem = i32;
impl PrecisionSettings for DoublePrecisionSettings {
    type FloatElem = f64;
    type IntElem = i64;
impl PrecisionSettings for HalfPrecisionSettings {
    type FloatElem = half::f16;
    type IntElem = i16;
/// Deserialize the value into [`TensorData`].
fn deserialize_data<'de, E, De>(deserializer: De) -> Result<TensorData, De::Error>
where
    E: Element + Deserialize<'de>,
    De: serde::Deserializer<'de>,
{
// do not convert quantized tensors
/// This struct implements serde to lazily serialize and deserialize a float tensor
/// using the given [record settings](RecordSettings).
pub struct FloatTensorSerde<S: PrecisionSettings> {
/// This struct implements serde to lazily serialize and deserialize an int tensor
/// using the given [record settings](RecordSettings).
pub struct IntTensorSerde<S: PrecisionSettings> {
/// This struct implements serde to lazily serialize and deserialize an bool tensor.
pub struct BoolTensorSerde {
// --- SERDE IMPLEMENTATIONS --- //
impl<S: PrecisionSettings> Serialize for FloatTensorSerde<S> {
    fn serialize<Se>(&self, serializer: Se) -> Result<Se::Ok, Se::Error>
    where
        Se: serde::Serializer,
    {
impl<'de, S: PrecisionSettings> Deserialize<'de> for FloatTensorSerde<S> {
    fn deserialize<De>(deserializer: De) -> Result<Self, De::Error>
    where
        De: serde::Deserializer<'de>,
    {
impl<S: PrecisionSettings> Serialize for IntTensorSerde<S> {
    fn serialize<Se>(&self, serializer: Se) -> Result<Se::Ok, Se::Error>
    where
        Se: serde::Serializer,
    {
impl<'de, S: PrecisionSettings> Deserialize<'de> for IntTensorSerde<S> {
    fn deserialize<De>(deserializer: De) -> Result<Self, De::Error>
    where
        De: serde::Deserializer<'de>,
    {
impl Serialize for BoolTensorSerde {
    fn serialize<Se>(&self, serializer: Se) -> Result<Se::Ok, Se::Error>
    where
        Se: serde::Serializer,
    {
impl<'de> Deserialize<'de> for BoolTensorSerde {
    fn deserialize<De>(deserializer: De) -> Result<Self, De::Error>
    where
        De: serde::Deserializer<'de>,
    {
// --- RECORD IMPLEMENTATIONS --- //
impl<B: Backend, const D: usize> Record<B> for Tensor<B, D> {
    type Item<S: PrecisionSettings> = FloatTensorSerde<S>;
    fn into_item<S: PrecisionSettings>(self) -> Self::Item<S> {
// do not convert quantized tensors
    fn from_item<S: PrecisionSettings>(item: Self::Item<S>, device: &B::Device) -> Self {
// do not convert quantized tensors
impl<B: Backend, const D: usize> Record<B> for Tensor<B, D, Int> {
    type Item<S: PrecisionSettings> = IntTensorSerde<S>;
    fn into_item<S: PrecisionSettings>(self) -> Self::Item<S> {
    fn from_item<S: PrecisionSettings>(item: Self::Item<S>, device: &B::Device) -> Self {
impl<B: Backend, const D: usize> Record<B> for Tensor<B, D, Bool> {
    type Item<S: PrecisionSettings> = BoolTensorSerde;
    fn into_item<S: PrecisionSettings>(self) -> Self::Item<S> {
    fn from_item<S: PrecisionSettings>(item: Self::Item<S>, device: &B::Device) -> Self {
pub struct TestEmptyStructConfig {
pub struct TestStructConfig {
pub enum TestEnumConfig {
fn file_path(file_name: &str) -> std::path::PathBuf {
fn struct_config_should_impl_serde() {
fn struct_config_should_impl_clone() {
fn struct_config_should_impl_display() {
fn enum_config_no_value_should_impl_serde() {
fn enum_config_one_value_should_impl_serde() {
fn enum_config_multiple_values_should_impl_serde() {
fn enum_config_should_impl_clone() {
fn enum_config_should_impl_display() {
fn struct_config_can_load_binary() {
pub type TestBackend = burn_ndarray::NdArray<f32>;
pub type TestAutodiffBackend = burn_autodiff::Autodiff<TestBackend>;
pub struct ModuleBasic<B: Backend> {
struct ModuleTensorConstInt<B: Backend> {
impl<B: Backend> ModuleBasic<B> {
    fn new(device: &B::Device) -> Self {
struct ModuleWithConstGeneric<B: Backend, const N: usize> {
struct ModuleWithGenericModule<B: Backend, M> {
enum ModuleEnum<B: Backend> {
enum ModuleEnumNested<B: Backend> {
enum ModuleEnumWithGenericModule<B: Backend, M: Module<B>> {
pub struct ModuleComposed<B: Backend> {
impl<B: Backend> ModuleComposed<B> {
    fn new(device: &B::Device) -> Self {
    type RecordItem<M, B, S> = <<M as Module<B>>::Record as Record<B>>::Item<S>;
    fn implements_clone<T: Clone>() {
    fn basic_implements_clone<B: Backend, S: PrecisionSettings>() {
    fn generic_implements_clone<B, S, M>()
    where
        B: Backend,
        S: PrecisionSettings,
        M: Module<B> + ModuleDisplay,
        RecordItem<M, B, S>: Clone,
    {
    fn should_load_from_record_basic() {
    fn should_load_from_record_compose() {
    fn should_load_from_record_enum() {
    fn should_load_from_record_const_generic() {
    fn should_panic_load_from_incorrect_enum_variant() {
    fn should_calculate_num_params_basic() {
    fn should_output_state_composed() {
    fn should_calculate_num_params_enum() {
    fn should_have_grad_by_default() {
    fn should_have_no_grad_after_no_grad() {
    fn should_have_grad_when_from_record() {
// Even when param is no_grad,
    fn calculate_grads(
        module: &ModuleBasic<TestAutodiffBackend>,
    ) -> <TestAutodiffBackend as AutodiffBackend>::Gradients {
// It compiles
pub struct TestWithBackendRecord<B: Backend> {
// It compiles
pub struct TestWithoutBackendRecord {
    type TestBackend = burn_ndarray::NdArray<f32>;
/// Simple linear module.
    pub struct Linear<B: Backend> {
    impl<B: Backend> Linear<B> {
        pub fn new(in_features: usize, out_features: usize, device: &B::Device) -> Self {
    pub struct Model<B: Backend> {
    pub struct ModelNewOptionalField<B: Backend> {
    pub struct ModelNewConstantField<B: Backend> {
    pub struct ModelNewFieldOrders<B: Backend> {
    fn deserialize_with_new_optional_field_works_with_default_file_recorder() {
    fn deserialize_with_removed_optional_field_works_with_default_file_recorder() {
    fn deserialize_with_new_constant_field_works_with_default_file_recorder() {
    fn deserialize_with_removed_constant_field_works_with_default_file_recorder() {
    fn deserialize_with_new_field_order_works_with_default_file_recorder() {
    fn deserialize_with_new_optional_field_works_with_pretty_json() {
    fn deserialize_with_removed_optional_field_works_with_pretty_json() {
    fn deserialize_with_new_constant_field_works_with_pretty_json() {
    fn deserialize_with_removed_constant_field_works_with_pretty_json() {
    fn deserialize_with_new_field_order_works_with_pretty_json() {
    fn deserialize_with_new_optional_field_doesnt_works_with_bin_file_recorder() {
    fn deserialize_with_removed_optional_field_works_with_bin_file_recorder() {
    fn deserialize_with_new_constant_field_works_with_bin_file_recorder() {
    fn deserialize_with_removed_constant_field_works_with_bin_file_recorder() {
    fn deserialize_with_new_field_order_works_with_bin_file_recorder() {
    fn file_path(filename: String) -> PathBuf {
    fn test_tensor_serde() {
    fn deserialize_with_new_optional_field<R>(name: &str, recorder: R) -> Result<(), RecorderError>
    where
        R: FileRecorder<TestBackend>,
    {
    fn deserialize_with_removed_optional_field<R>(
        name: &str,
        recorder: R,
    ) -> Result<(), RecorderError>
    where
        R: FileRecorder<TestBackend>,
    {
    fn deserialize_with_new_constant_field<R>(name: &str, recorder: R) -> Result<(), RecorderError>
    where
        R: FileRecorder<TestBackend>,
    {
    fn deserialize_with_removed_constant_field<R>(
        name: &str,
        recorder: R,
    ) -> Result<(), RecorderError>
    where
        R: FileRecorder<TestBackend>,
    {
    fn deserialize_with_new_field_order<R>(name: &str, recorder: R) -> Result<(), RecorderError>
    where
        R: FileRecorder<TestBackend>,
    {
pub type Cuda<F = f32, I = i32> = CubeBackend<CudaRuntime, F, I, u8>;
pub type Cuda<F = f32, I = i32> = burn_fusion::Fusion<CubeBackend<CudaRuntime, F, I, u8>>;
    pub type TestRuntime = cubecl::cuda::CudaRuntime;
// TODO: Add tests for bf16
// burn_cubecl::testgen_all!([f16, f32], [i8, i16, i32, i64], [u8, u32]);
/// A tensor representation containing a reference to a tensor resource with a given shape.
pub struct TensorHandle<H: Clone> {
/// The type that can be used to point to a tensor of any kind.
/// The shape associated to the tensor.
/// Backend extension trait that allows an existing [backend](Backend) to use the Burn tensor
/// intermediate representation for compilation purpose or other...
pub trait BackendIr: Backend {
/// The type that can be used to point to a tensor of any kind.
    type Handle: Sync + Send + Clone;
/// Convert a [handle](BackendIr::Handle) to a [float tensor](Backend::FloatTensorPrimitive).
    fn float_tensor(handle: TensorHandle<Self::Handle>) -> FloatTensor<Self>;
/// Convert a [handle](BackendIr::Handle) to an [int tensor](Backend::IntTensorPrimitive).
    fn int_tensor(handle: TensorHandle<Self::Handle>) -> IntTensor<Self>;
/// Convert a [handle](BackendIr::Handle) to a [bool tensor](Backend::BoolTensorPrimitive).
    fn bool_tensor(handle: TensorHandle<Self::Handle>) -> BoolTensor<Self>;
/// Convert a [handle](BackendIr::Handle) to a [quantized tensor](Backend::QuantizedTensorPrimitive).
    fn quantized_tensor(handle: TensorHandle<Self::Handle>) -> QuantizedTensor<Self>;
/// Convert a [float tensor](Backend::FloatTensorPrimitive) to a [handle](BackendIr::Handle).
    fn float_tensor_handle(tensor: FloatTensor<Self>) -> Self::Handle;
/// Convert an [int tensor](Backend::IntTensorPrimitive) to a [handle](BackendIr::Handle).
    fn int_tensor_handle(tensor: IntTensor<Self>) -> Self::Handle;
/// Convert a [bool tensor](Backend::BoolTensorPrimitive) to a [handle](BackendIr::Handle).
    fn bool_tensor_handle(tensor: BoolTensor<Self>) -> Self::Handle;
/// Convert a [quantized tensor](Backend::QuantizedTensorPrimitive) to a [handle](BackendIr::Handle).
    fn quantized_tensor_handle(tensor: QuantizedTensor<Self>) -> Self::Handle;
/// Handle which points to a backend tensor primitive kind.
pub enum HandleKind<B: Backend> {
/// Float tensor handle.
/// Int tensor handle.
/// Bool tensor handle.
/// Quantized tensor handle.
impl<B: Backend> HandleKind<B> {
/// Returns the handle kind name.
    pub fn name(&self) -> &str {
/// Keep all [tensor handles](BackendIr::Handle) in one place and ensure that all resources
/// are used optimally.
pub struct HandleContainer<H> {
impl<H: Clone> HandleContainer<H> {
/// Fork the container, useful for autotune.
    pub fn fork(&self) -> Self {
impl<H> core::fmt::Debug for HandleContainer<H> {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
// only care about the IDs when debugging
/// Backend [tensor handle](BackendIr::Handle) wrapper tracking their creation state
pub enum Handle<H> {
/// No [tensor handle](BackendIr::Handle) has been created yet
/// A [tensor handle](BackendIr::Handle) has been created
impl<H: Clone> HandleContainer<H> {
/// Create a new HandleContainer
    pub fn new() -> Self {
/// Register a handle for the given [tensor id](TensorId).
    pub fn register_handle(&mut self, id: TensorId, handle: H) {
/// Whether an handle exists.
    pub fn has_handle(&mut self, id: &TensorId) -> bool {
/// Get the handle for the given [tensor id](TensorId). The status is used to determine if the
/// tensor should be popped out of the current tensor map, necessary for inplace operations.
///
/// # Warnings
///
/// Make sure the status corresponds to the operation you want to execute the handle on,
/// otherwise you might remove a tensor handle that will be required in the future.
    pub fn get_handle(&mut self, id: &TensorId, status: &TensorStatus) -> H {
/// Get the tensor handle for the given [tensor intermediate representation](TensorIr).
    pub fn get_tensor_handle(&mut self, tensor: &TensorIr) -> TensorHandle<H> {
/// Get the [float tensor](burn_tensor::backend::Backend::FloatTensorPrimitive) corresponding to the
/// given [tensor intermediate representation](TensorIr).
    pub fn get_float_tensor<B>(&mut self, tensor: &TensorIr) -> B::FloatTensorPrimitive
    where
        B: BackendIr<Handle = H>,
    {
/// Get the [int tensor](burn_tensor::backend::Backend::IntTensorPrimitive) corresponding to the
/// given [tensor intermediate representation](TensorIr).
    pub fn get_int_tensor<B>(&mut self, tensor: &TensorIr) -> B::IntTensorPrimitive
    where
        B: BackendIr<Handle = H>,
    {
/// Get the [bool tensor](burn_tensor::backend::Backend::BoolTensorPrimitive) corresponding to the
/// given [tensor intermediate representation](TensorIr).
    pub fn get_bool_tensor<B>(&mut self, tensor: &TensorIr) -> B::BoolTensorPrimitive
    where
        B: BackendIr<Handle = H>,
    {
/// Get the [quantized tensor](burn_tensor::backend::Backend::QuantizedTensorPrimitive) corresponding to the
/// given [tensor intermediate representation](TensorIr).
    pub fn get_quantized_tensor<B>(&mut self, tensor: &TensorIr) -> B::QuantizedTensorPrimitive
    where
        B: BackendIr<Handle = H>,
    {
/// Register a new [float tensor](burn_tensor::backend::Backend::FloatTensorPrimitive) with the corresponding [tensor id](TensorId).
    pub fn register_float_tensor<B>(&mut self, id: &TensorId, tensor: B::FloatTensorPrimitive)
    where
        B: BackendIr<Handle = H>,
    {
/// Register a new [quantized tensor](burn_tensor::backend::Backend::QuantizedTensorPrimitive) with the corresponding [tensor ids](TensorId).
    pub fn register_quantized_tensor<B>(
        &mut self,
        id: &TensorId,
        tensor: B::QuantizedTensorPrimitive,
    ) where
        B: BackendIr<Handle = H>,
    {
/// Register a new [int tensor](burn_tensor::backend::Backend::IntTensorPrimitive) with the corresponding [tensor id](TensorId).
    pub fn register_int_tensor<B>(&mut self, id: &TensorId, tensor: B::IntTensorPrimitive)
    where
        B: BackendIr<Handle = H>,
    {
/// Register a new [bool tensor](burn_tensor::backend::Backend::BoolTensorPrimitive) with the corresponding [tensor id](TensorId).
    pub fn register_bool_tensor<B>(&mut self, id: &TensorId, tensor: B::BoolTensorPrimitive)
    where
        B: BackendIr<Handle = H>,
    {
/// Lazily create a new empty tensor and return its corresponding [tensor id](TensorId).
    pub fn create_tensor_uninit(&mut self) -> TensorId {
/// Remove tensor handle from container.
    pub fn remove_handle(&mut self, id: TensorId) -> Option<Handle<H>> {
/// Remove tensor handle from container if writable
    pub fn free(&mut self, tensor: &TensorIr) {
/// Returns the number of handles.
    pub fn num_handles(&self) -> usize {
//! Burn intermediate representation.
/// Custom operation in fusion stream, declaring its inputs and outputs.
pub struct CustomOpIr {
/// Unique identifier of the operation.
/// Input tensors used in the custom operation.
/// Output tensors used in the custom operation.
impl CustomOpIr {
/// Create a new custom operation intermediate representation.
    pub fn new(id: &'static str, inputs: &[TensorIr], outputs: &[TensorIr]) -> Self {
/// Cast the intermediate representation, and get the in and output tensors.
    pub fn as_fixed<const N_IN: usize, const N_OUT: usize>(
        &self,
    ) -> (&[TensorIr;
    fn nodes(&self) -> Vec<&TensorIr> {
/// irribe all tensor operations possible.
pub enum OperationIr {
/// Basic operation on a float tensor.
/// Basic operation on an int tensor.
/// Basic operation on a bool tensor.
/// Numeric operation on a float tensor.
/// Numeric operation on an int tensor.
/// Operation specific to a bool tensor.
/// Operation specific to an int tensor.
/// Operation specific to a float tensor.
/// Module operation.
/// Initialize operation.
/// A custom operation.
/// A tensor is dropped.
/// Operation intermediate representation specific to a float tensor.
pub enum FloatOperationIr {
/// Operation corresponding to [exp](burn_tensor::ops::FloatTensorOps::float_exp).
/// Operation corresponding to [log](burn_tensor::ops::FloatTensorOps::float_log).
/// Operation corresponding to [log1p](burn_tensor::ops::FloatTensorOps::float_log1p).
/// Operation corresponding to [erf](burn_tensor::ops::FloatTensorOps::float_erf).
/// Operation corresponding to [powf_scalar](burn_tensor::ops::FloatTensorOps::float_powf_scalar).
/// Operation corresponding to [sqrt](burn_tensor::ops::FloatTensorOps::float_sqrt).
/// Operation corresponding to [cos](burn_tensor::ops::FloatTensorOps::float_cos).
/// Operation corresponding to [sin](burn_tensor::ops::FloatTensorOps::float_sin).
/// Operation corresponding to [tanh](burn_tensor::ops::FloatTensorOps::float_tanh).
/// Operation corresponding to [round](burn_tensor::ops::FloatTensorOps::float_round).
/// Operation corresponding to [floor](burn_tensor::ops::FloatTensorOps::float_floor).
/// Operation corresponding to [ceil](burn_tensor::ops::FloatTensorOps::float_ceil).
/// Operation corresponding to [into_int](burn_tensor::ops::FloatTensorOps::float_into_int).
/// Operation corresponding to [matmul](burn_tensor::ops::FloatTensorOps::float_matmul).
/// Operation corresponding to [random](burn_tensor::ops::FloatTensorOps::float_random).
/// Operation corresponding to [recip](burn_tensor::ops::FloatTensorOps::float_recip).
/// Operation corresponding to [quantize](burn_tensor::ops::QTensorOps::quantize).
/// Operation corresponding to [dequantize](burn_tensor::ops::QTensorOps::dequantize).
/// Operation intermediate representation specific to module.
pub enum ModuleOperationIr {
/// Operation corresponding to [embedding](burn_tensor::ops::ModuleOps::embedding).
/// Operation corresponding to [embedding_backward](burn_tensor::ops::ModuleOps::embedding_backward).
/// Operation corresponding to [conv1d](burn_tensor::ops::ModuleOps::conv1d).
/// Operation corresponding to [conv2d](burn_tensor::ops::ModuleOps::conv2d).
/// Operation corresponding to [conv3d](burn_tensor::ops::ModuleOps::conv3d).
/// Operation corresponding to [deform_conv2d](burn_tensor::ops::ModuleOps::deform_conv2d)
/// Operation corresponding to [deform_conv2d_backward](burn_tensor::ops::ModuleOps::deform_conv2d_backward)
/// Operation corresponding to [conv transpose 1d](burn_tensor::ops::ModuleOps::conv_transpose1d).
/// Operation corresponding to [conv transpose 2d](burn_tensor::ops::ModuleOps::conv_transpose2d).
/// Operation corresponding to [conv transpose 3d](burn_tensor::ops::ModuleOps::conv_transpose3d).
/// Operation corresponding to [avg pool 1d](burn_tensor::ops::ModuleOps::avg_pool1d).
/// Operation corresponding to [avg pool 2d](burn_tensor::ops::ModuleOps::avg_pool2d).
/// Operation corresponding to
/// [avg pool 1d backward](burn_tensor::ops::ModuleOps::avg_pool1d_backward).
/// Operation corresponding to
/// [avg pool 2d backward](burn_tensor::ops::ModuleOps::avg_pool2d_backward).
/// Operation corresponding to
/// [adaptive avg pool 1d](burn_tensor::ops::ModuleOps::adaptive_avg_pool1d).
/// Operation corresponding to
/// [adaptive avg pool 2d](burn_tensor::ops::ModuleOps::adaptive_avg_pool2d).
/// Operation corresponding to
/// [adaptive avg pool 1d backward](burn_tensor::ops::ModuleOps::adaptive_avg_pool1d_backward).
/// Operation corresponding to
/// [adaptive avg pool 2d backward](burn_tensor::ops::ModuleOps::adaptive_avg_pool2d_backward).
/// Operation corresponding to
/// [max pool 1d](burn_tensor::ops::ModuleOps::max_pool1d).
/// Operation corresponding to
/// [max pool 1d with indices](burn_tensor::ops::ModuleOps::max_pool1d_with_indices).
/// Operation corresponding to
/// [max pool 1d with indices backward](burn_tensor::ops::ModuleOps::max_pool1d_with_indices_backward).
/// Operation corresponding to
/// [max pool 2d](burn_tensor::ops::ModuleOps::max_pool1d).
/// Operation corresponding to
/// [max pool 2d with indices](burn_tensor::ops::ModuleOps::max_pool2d_with_indices).
/// Operation corresponding to
/// [max pool 2d with indices backward](burn_tensor::ops::ModuleOps::max_pool2d_with_indices_backward).
/// Operation corresponding to [interpolate](burn_tensor::ops::ModuleOps::interpolate).
/// Operation corresponding to [interpolate backward](burn_tensor::ops::ModuleOps::interpolate_backward).
/// Basic operations that can be done on any tensor type.
pub enum BaseOperationIr {
/// Operation corresponding to:
///
/// Float => [to device](burn_tensor::ops::FloatTensorOps::float_to_device).
/// Int => [to device](burn_tensor::ops::IntTensorOps::int_to_device).
/// Bool => [to device](burn_tensor::ops::BoolTensorOps::bool_to_device).
/// Operation corresponding to:
///
/// Float => [reshape](burn_tensor::ops::FloatTensorOps::float_reshape).
/// Int => [reshape](burn_tensor::ops::IntTensorOps::int_reshape).
/// Bool => [reshape](burn_tensor::ops::BoolTensorOps::bool_reshape).
/// Operation corresponding to:
///
/// Float => [swap_dims](burn_tensor::ops::FloatTensorOps::float_swap_dims).
/// Int => [swap_dims](burn_tensor::ops::IntTensorOps::int_swap_dims).
/// Bool => [swap_dims](burn_tensor::ops::BoolTensorOps::bool_swap_dims).
/// Operation corresponding to:
///
/// Float => [permute](burn_tensor::ops::FloatTensorOps::float_permute).
/// Int => [permute](burn_tensor::ops::IntTensorOps::int_permute).
/// Bool => [permute](burn_tensor::ops::BoolTensorOps::bool_permute).
/// Operation corresponding to:
/// Float => [flip](burn_tensor::ops::FloatTensorOps::float_flip).
/// Int => [flip](burn_tensor::ops::IntTensorOps::int_flip).
/// Bool => [flip](burn_tensor::ops::BoolTensorOps::bool_flip).
/// Operation corresponding to:
///
/// Float => [expand](burn_tensor::ops::FloatTensorOps::float_expand).
/// Int => [expand](burn_tensor::ops::IntTensorOps::int_expand).
/// Bool => [expand](burn_tensor::ops::BoolTensorOps::bool_expand).
/// Operation corresponding to:
///
/// Float => [slice](burn_tensor::ops::FloatTensorOps::float_slice).
/// Int => [slice](burn_tensor::ops::IntTensorOps::int_slice).
/// Bool => [slice](burn_tensor::ops::BoolTensorOps::bool_slice).
/// Operation corresponding to:
///
/// Float => [slice assign](burn_tensor::ops::FloatTensorOps::float_slice_assign).
/// Int => [slice assign](burn_tensor::ops::IntTensorOps::int_slice_assign).
/// Bool => [slice assign](burn_tensor::ops::BoolTensorOps::bool_slice_assign).
/// Operation corresponding to:
///
/// Float => [equal](burn_tensor::ops::FloatTensorOps::float_equal).
/// Int => [equal](burn_tensor::ops::IntTensorOps::int_equal).
/// Bool => [equal](burn_tensor::ops::BoolTensorOps::bool_equal).
/// Operation corresponding to:
///
/// Float => [repeat dim](burn_tensor::ops::FloatTensorOps::float_repeat_dim).
/// Int => [repeat dim](burn_tensor::ops::IntTensorOps::int_repeat_dim).
/// Bool => [repeat dim](burn_tensor::ops::BoolTensorOps::bool_repeat_dim).
/// Operation corresponding to:
///
/// Float => [cat](burn_tensor::ops::FloatTensorOps::float_cat).
/// Int => [cat](burn_tensor::ops::IntTensorOps::int_cat).
/// Bool => [cat](burn_tensor::ops::BoolTensorOps::bool_cat).
/// Cast operation, no direct operation and should be supported by fusion backend.
/// Operation corresponding to:
///
/// Float => [empty](burn_tensor::ops::FloatTensorOps::float_empty).
/// Int => [empty](burn_tensor::ops::IntTensorOps::int_empty).
/// Bool => [empty](burn_tensor::ops::BoolTensorOps::bool_empty).
/// Numeric operations on int and float tensors.
pub enum NumericOperationIr<E> {
/// Operation corresponding to:
///
/// Float => [add](burn_tensor::ops::FloatTensorOps::float_add).
/// Int => [add](burn_tensor::ops::IntTensorOps::int_add).
/// Operation corresponding to:
///
/// Float => [add scalar](burn_tensor::ops::FloatTensorOps::float_add_scalar).
/// Int => [add scalar](burn_tensor::ops::IntTensorOps::int_add_scalar).
/// Operation corresponding to:
///
/// Float => [sub](burn_tensor::ops::FloatTensorOps::float_sub).
/// Int => [sub](burn_tensor::ops::IntTensorOps::int_sub).
/// Operation corresponding to:
///
/// Float => [sub scalar](burn_tensor::ops::FloatTensorOps::float_sub_scalar).
/// Int => [sub scalar](burn_tensor::ops::IntTensorOps::int_sub_scalar).
/// Operation corresponding to:
///
/// Float => [div](burn_tensor::ops::FloatTensorOps::float_div).
/// Int => [div](burn_tensor::ops::IntTensorOps::int_div).
/// Operation corresponding to:
///
/// Float => [div scalar](burn_tensor::ops::FloatTensorOps::float_div_scalar).
/// Int => [div scalar](burn_tensor::ops::IntTensorOps::int_div_scalar).
/// Operation corresponding to:
///
/// Float => [rem](burn_tensor::ops::FloatTensorOps::float_remainder).
/// Int => [rem](burn_tensor::ops::IntTensorOps::int_remainder).
/// Operation corresponding to:
///
/// Float => [rem scalar](burn_tensor::ops::FloatTensorOps::float_remainder_scalar).
/// Int => [rem scalar](burn_tensor::ops::IntTensorOps::int_remainder_scalar).
/// Operation corresponding to:
///
/// Float => [mul](burn_tensor::ops::FloatTensorOps::float_mul).
/// Int => [mul](burn_tensor::ops::IntTensorOps::int_mul).
/// Operation corresponding to:
///
/// Float => [mul scalar](burn_tensor::ops::FloatTensorOps::float_mul_scalar).
/// Int => [mul scalar](burn_tensor::ops::IntTensorOps::int_mul_scalar).
/// Operation corresponding to:
///
/// Float => [abs](burn_tensor::ops::FloatTensorOps::float_abs).
/// Int => [abs](burn_tensor::ops::IntTensorOps::int_abs).
/// Operation corresponding to:
///
/// Float => [ones](burn_tensor::ops::FloatTensorOps::float_ones).
/// Int => [ones](burn_tensor::ops::IntTensorOps::int_ones).
/// Operation corresponding to:
///
/// Float => [zeros](burn_tensor::ops::FloatTensorOps::float_zeros).
/// Int => [zeros](burn_tensor::ops::IntTensorOps::int_zeros).
/// Operation corresponding to:
///
/// Float => [full](burn_tensor::ops::FloatTensorOps::float_full).
/// Int => [full](burn_tensor::ops::IntTensorOps::int_full).
/// Operation corresponding to:
///
/// Float => [gather](burn_tensor::ops::FloatTensorOps::float_gather).
/// Int => [gather](burn_tensor::ops::IntTensorOps::int_gather).
/// Operation corresponding to:
///
/// Float => [scatter](burn_tensor::ops::FloatTensorOps::float_scatter).
/// Int => [scatter](burn_tensor::ops::IntTensorOps::int_scatter).
/// Operation corresponding to:
///
/// Float => [select](burn_tensor::ops::FloatTensorOps::float_select).
/// Int => [select](burn_tensor::ops::IntTensorOps::int_select).
/// Operation corresponding to:
///
/// Float => [select assign](burn_tensor::ops::FloatTensorOps::float_select_assign).
/// Int => [select assign](burn_tensor::ops::IntTensorOps::int_select_assign).
/// Operation corresponding to:
///
/// Float => [mask where](burn_tensor::ops::FloatTensorOps::float_mask_where).
/// Int => [mask where](burn_tensor::ops::IntTensorOps::int_mask_where).
/// Operation corresponding to:
///
/// Float => [mask fill](burn_tensor::ops::FloatTensorOps::float_mask_fill).
/// Int => [mask fill](burn_tensor::ops::IntTensorOps::int_mask_fill).
/// Operation corresponding to:
///
/// Float => [mean dim](burn_tensor::ops::FloatTensorOps::float_mean_dim).
/// Int => [mean dim](burn_tensor::ops::IntTensorOps::int_mean_dim).
/// Operation corresponding to:
///
/// Float => [mean](burn_tensor::ops::FloatTensorOps::float_mean).
/// Int => [mean](burn_tensor::ops::IntTensorOps::int_mean).
/// Operation corresponding to:
///
/// Float => [sum](burn_tensor::ops::FloatTensorOps::float_sum).
/// Int => [sum](burn_tensor::ops::IntTensorOps::int_sum).
/// Operation corresponding to:
///
/// Float => [sum dim](burn_tensor::ops::FloatTensorOps::float_sum_dim).
/// Int => [sum dim](burn_tensor::ops::IntTensorOps::int_sum_dim).
/// Operation corresponding to:
///
/// Float => [prod](burn_tensor::ops::FloatTensorOps::float_prod).
/// Int => [prod](burn_tensor::ops::IntTensorOps::int_prod).
/// Operation corresponding to:
///
/// Float => [prod dim](burn_tensor::ops::FloatTensorOps::float_prod_dim).
/// Int => [prod dim](burn_tensor::ops::IntTensorOps::int_prod_dim).
/// Operation corresponding to:
///
/// Float => [equal elem](burn_tensor::ops::FloatTensorOps::float_equal_elem).
/// Int => [equal elem](burn_tensor::ops::IntTensorOps::int_equal_elem).
/// Operation corresponding to:
///
/// Float => [greater](burn_tensor::ops::FloatTensorOps::float_greater).
/// Int => [greater](burn_tensor::ops::IntTensorOps::int_greater).
/// Operation corresponding to:
///
/// Float => [greater elem](burn_tensor::ops::FloatTensorOps::float_greater_elem).
/// Int => [greater elem](burn_tensor::ops::IntTensorOps::int_greater_elem).
/// Operation corresponding to:
///
/// Float => [greater equal](burn_tensor::ops::FloatTensorOps::float_greater_elem).
/// Int => [greater elem](burn_tensor::ops::IntTensorOps::int_greater_elem).
/// Operation corresponding to:
///
/// Float => [greater equal elem](burn_tensor::ops::FloatTensorOps::float_greater_equal_elem).
/// Int => [greater equal elem](burn_tensor::ops::IntTensorOps::int_greater_equal_elem).
/// Operation corresponding to:
///
/// Float => [lower](burn_tensor::ops::FloatTensorOps::float_lower).
/// Int => [lower](burn_tensor::ops::IntTensorOps::int_lower).
/// Operation corresponding to:
///
/// Float => [lower elem](burn_tensor::ops::FloatTensorOps::float_lower_elem).
/// Int => [lower elem](burn_tensor::ops::IntTensorOps::int_lower_elem).
/// Operation corresponding to:
///
/// Float => [lower equal](burn_tensor::ops::FloatTensorOps::float_lower_equal).
/// Int => [lower equal](burn_tensor::ops::IntTensorOps::int_lower_equal).
/// Operation corresponding to:
///
/// Float => [lower equal elem](burn_tensor::ops::FloatTensorOps::float_lower_equal_elem).
/// Int => [lower equal elem](burn_tensor::ops::IntTensorOps::int_lower_equal_elem).
/// Operation corresponding to:
///
/// Float => [argmax](burn_tensor::ops::FloatTensorOps::float_argmax).
/// Int => [argmax](burn_tensor::ops::IntTensorOps::int_argmax).
/// Operation corresponding to:
///
/// Float => [argmin](burn_tensor::ops::FloatTensorOps::float_argmin).
/// Int => [argmin](burn_tensor::ops::IntTensorOps::int_argmin).
/// Operation corresponding to:
///
/// Float => [max](burn_tensor::ops::FloatTensorOps::float_max).
/// Int => [max](burn_tensor::ops::IntTensorOps::int_max).
/// Operation corresponding to:
///
/// Float => [max dim with indices](burn_tensor::ops::FloatTensorOps::float_max_dim_with_indices).
/// Int => [max dim with indices](burn_tensor::ops::IntTensorOps::int_max_dim_with_indices).
/// Operation corresponding to:
///
/// Float => [min dim with indices](burn_tensor::ops::FloatTensorOps::float_min_dim_with_indices).
/// Int => [min dim with indices](burn_tensor::ops::IntTensorOps::int_min_dim_with_indices).
/// Operation corresponding to:
///
/// Float => [min](burn_tensor::ops::FloatTensorOps::float_min).
/// Int => [min](burn_tensor::ops::IntTensorOps::int_min).
/// Operation corresponding to:
///
/// Float => [max dim](burn_tensor::ops::FloatTensorOps::float_max_dim).
/// Int => [max dim](burn_tensor::ops::IntTensorOps::int_max_dim).
/// Operation corresponding to:
///
/// Float => [min dim](burn_tensor::ops::FloatTensorOps::float_min_dim).
/// Int => [min dim](burn_tensor::ops::IntTensorOps::int_min_dim).
/// Operation corresponding to:
///
/// Float => [max_abs](burn_tensor::ops::FloatTensorOps::float_max_abs).
/// Int => [max_abs](burn_tensor::ops::IntTensorOps::int_max_abs).
/// Operation corresponding to:
///
/// Float => [max_abs dim](burn_tensor::ops::FloatTensorOps::float_max_abs_dim).
/// Int => [max_abs dim](burn_tensor::ops::IntTensorOps::int_max_abs_dim).
/// Operation corresponding to:
///
/// Float => [clamp](burn_tensor::ops::FloatTensorOps::float_clamp).
/// Int => [clamp](burn_tensor::ops::IntTensorOps::int_clamp).
/// Operation corresponding to:
///
/// Int => [random](burn_tensor::ops::IntTensorOps::int_random).
/// Operation corresponding to:
///
/// Float => [powf](burn_tensor::ops::FloatTensorOps::float_powf).
/// Int => [powf](burn_tensor::ops::IntTensorOps::int_powf).
/// Operation intermediate representation specific to an int tensor.
pub enum IntOperationIr {
/// Operation corresponding to [into float](burn_tensor::ops::IntTensorOps::int_into_float).
/// Operation corresponding to:
///
/// Int => [bitwise and](burn_tensor::ops::IntTensorOps::bitwise_and).
/// Operation corresponding to:
///
/// Int => [bitwise and scalar](burn_tensor::ops::IntTensorOps::bitwise_and_scalar).
/// Operation corresponding to:
///
/// Int => [bitwise or](burn_tensor::ops::IntTensorOps::bitwise_or).
/// Operation corresponding to:
///
/// Int => [bitwise or scalar](burn_tensor::ops::IntTensorOps::bitwise_or_scalar).
/// Operation corresponding to:
///
/// Int => [bitwise xor](burn_tensor::ops::IntTensorOps::bitwise_xor).
/// Operation corresponding to:
///
/// Int => [bitwise xor scalar](burn_tensor::ops::IntTensorOps::bitwise_xor_scalar).
/// Operation corresponding to:
///
/// Int => [bitwise not](burn_tensor::ops::IntTensorOps::bitwise_not).
/// Operation corresponding to:
///
/// Int => [bitwise left shift](burn_tensor::ops::IntTensorOps::bitwise_left_shift).
/// Operation corresponding to:
///
/// Int => [bitwise left shift scalar](burn_tensor::ops::IntTensorOps::bitwise_left_shift_scalar).
/// Operation corresponding to:
///
/// Int => [bitwise right shift](burn_tensor::ops::IntTensorOps::bitwise_right_shift).
/// Operation corresponding to:
///
/// Int => [bitwise right shift scalar](burn_tensor::ops::IntTensorOps::bitwise_right_shift_scalar).
/// Operation intermediate representation specific to a bool tensor.
pub enum BoolOperationIr {
/// Operation corresponding to [into float](burn_tensor::ops::BoolTensorOps::bool_into_float).
/// Operation corresponding to [into int](burn_tensor::ops::BoolTensorOps::bool_into_int).
/// Operation corresponding to [not](burn_tensor::ops::BoolTensorOps::bool_not).
/// Operation corresponding to [and](burn_tensor::ops::BoolTensorOps::bool_and).
/// Operation corresponding to [or](burn_tensor::ops::BoolTensorOps::bool_or).
/// Swap dim operation intermediate representation.
pub struct SwapDimsOpIr {
/// Input tensor intermediate representation.
/// Output tensor intermediate representation.
/// The first dim to swap.
/// The second dim to swap.
/// Permute operation intermediate representation.
pub struct PermuteOpIr {
/// Input tensor intermediate representation.
/// Output tensor intermediate representation.
/// The new order of the dimensions.
/// Expand operation intermediate representation.
pub struct ExpandOpIr {
/// Input tensor intermediate representation.
/// Output tensor intermediate representation.
/// The new shape.
/// Flip operation intermediate representation.
pub struct FlipOpIr {
/// Input tensor intermediate representation.
/// Output tensor intermediate representation.
/// The dimensions to flip.
pub struct RandomOpIr {
/// Declares a tensor has been initialized.
///
/// It is necessary to register for proper orphan detection and avoid memory leak.
pub struct InitOperationIr {
/// The initialized tensor.
pub struct BinaryOpIr {
pub struct UnaryOpIr {
pub struct ScalarOpIr<E> {
// TODO: Make that an enum with `Value` and `Id` variants for relative/global
// conversion.
pub struct ReduceDimOpIr {
pub struct GatherOpIr {
pub struct ScatterOpIr {
pub struct SelectOpIr {
pub struct SelectAssignOpIr {
pub struct SliceOpIr {
pub struct SliceAssignOpIr {
pub struct MaskWhereOpIr {
pub struct MaskFillOpIr<E> {
pub struct ClampOpIr<E> {
pub struct RepeatDimOpIr {
pub struct CatOpIr {
pub struct ReduceDimWithIndicesOpIr {
pub struct EmbeddingOpIr {
pub struct EmbeddingBackwardOpIr {
pub struct Conv1dOpIr {
pub struct Conv2dOpIr {
pub struct DeformConv2dOpIr {
pub struct DeformConv2dBackwardOpIr {
pub struct Conv3dOpIr {
pub struct ConvTranspose1dOpIr {
pub struct ConvTranspose2dOpIr {
pub struct ConvTranspose3dOpIr {
pub struct Conv1dOptionsIr {
pub struct Conv2dOptionsIr {
pub struct DeformableConv2dOptionsIr {
pub struct Conv3dOptionsIr {
pub struct ConvTranspose1dOptionsIr {
pub struct ConvTranspose2dOptionsIr {
pub struct ConvTranspose3dOptionsIr {
/// Quantization parameters intermediate representation.
pub struct QuantizationParametersIr {
/// The scaling factor.
/// The zero-point offset.
pub struct QuantizeOpIr {
pub struct DequantizeOpIr {
impl From<ConvOptions<1>> for Conv1dOptionsIr {
    fn from(value: ConvOptions<1>) -> Self {
impl From<ConvOptions<2>> for Conv2dOptionsIr {
    fn from(value: ConvOptions<2>) -> Self {
impl From<ConvOptions<3>> for Conv3dOptionsIr {
    fn from(value: ConvOptions<3>) -> Self {
impl From<DeformConvOptions<2>> for DeformableConv2dOptionsIr {
    fn from(value: DeformConvOptions<2>) -> Self {
impl From<ConvTransposeOptions<1>> for ConvTranspose1dOptionsIr {
    fn from(value: ConvTransposeOptions<1>) -> Self {
impl From<ConvTransposeOptions<2>> for ConvTranspose2dOptionsIr {
    fn from(value: ConvTransposeOptions<2>) -> Self {
impl From<ConvTransposeOptions<3>> for ConvTranspose3dOptionsIr {
    fn from(value: ConvTransposeOptions<3>) -> Self {
impl From<Conv1dOptionsIr> for ConvOptions<1> {
    fn from(val: Conv1dOptionsIr) -> Self {
impl From<Conv2dOptionsIr> for ConvOptions<2> {
    fn from(val: Conv2dOptionsIr) -> Self {
impl From<Conv3dOptionsIr> for ConvOptions<3> {
    fn from(val: Conv3dOptionsIr) -> Self {
impl From<DeformableConv2dOptionsIr> for DeformConvOptions<2> {
    fn from(value: DeformableConv2dOptionsIr) -> Self {
impl From<ConvTranspose1dOptionsIr> for ConvTransposeOptions<1> {
    fn from(val: ConvTranspose1dOptionsIr) -> Self {
impl From<ConvTranspose2dOptionsIr> for ConvTransposeOptions<2> {
    fn from(val: ConvTranspose2dOptionsIr) -> Self {
impl From<ConvTranspose3dOptionsIr> for ConvTransposeOptions<3> {
    fn from(val: ConvTranspose3dOptionsIr) -> Self {
pub struct AvgPool1dOpIr {
pub struct AvgPool2dOpIr {
pub struct AvgPool1dBackwardOpIr {
pub struct AvgPool2dBackwardOpIr {
pub struct AdaptiveAvgPool1dOpIr {
pub struct AdaptiveAvgPool2dOpIr {
pub struct AdaptiveAvgPool1dBackwardOpIr {
pub struct AdaptiveAvgPool2dBackwardOpIr {
pub struct MaxPool1dOpIr {
pub struct MaxPool1dWithIndicesOpIr {
pub struct MaxPool1dWithIndicesBackwardOpIr {
pub struct MaxPool2dOpIr {
pub struct MaxPool2dWithIndicesOpIr {
pub struct MaxPool2dWithIndicesBackwardOpIr {
pub enum InterpolateModeIr {
pub struct InterpolateOptionsIr {
pub struct InterpolateOpIr {
impl From<InterpolateModeIr> for InterpolateMode {
    fn from(val: InterpolateModeIr) -> Self {
impl From<InterpolateOptionsIr> for InterpolateOptions {
    fn from(val: InterpolateOptionsIr) -> Self {
impl From<InterpolateMode> for InterpolateModeIr {
    fn from(val: InterpolateMode) -> Self {
impl From<InterpolateOptions> for InterpolateOptionsIr {
    fn from(val: InterpolateOptions) -> Self {
pub struct InterpolateBackwardOpIr {
impl OperationIr {
/// Get all [tensor](TensorIr) involved with the current operation.
    pub fn nodes(&self) -> Vec<&TensorIr> {
/// Set the given nodes that are [read write](super::TensorStatus::ReadWrite) to
/// [read only](super::TensorStatus::ReadOnly) in the current operation.
///
/// Returns the tensor that were updated with their original representation.
    pub fn mark_read_only(&mut self, nodes: &[TensorId]) -> Vec<TensorIr> {
impl BaseOperationIr {
    fn nodes(&self) -> Vec<&TensorIr> {
    fn mark_read_only(&mut self, nodes: &[TensorId]) -> Vec<TensorIr> {
impl<E: Element> NumericOperationIr<E> {
    fn nodes(&self) -> Vec<&TensorIr> {
    fn mark_read_only(&mut self, nodes: &[TensorId]) -> Vec<TensorIr> {
impl FloatOperationIr {
    fn nodes(&self) -> Vec<&TensorIr> {
    fn mark_read_only(&mut self, nodes: &[TensorId]) -> Vec<TensorIr> {
impl IntOperationIr {
    fn nodes(&self) -> Vec<&TensorIr> {
    fn mark_read_only(&mut self, nodes: &[TensorId]) -> Vec<TensorIr> {
impl BoolOperationIr {
    fn nodes(&self) -> Vec<&TensorIr> {
    fn mark_read_only(&mut self, nodes: &[TensorId]) -> Vec<TensorIr> {
impl ModuleOperationIr {
    fn nodes(&self) -> Vec<&TensorIr> {
    fn mark_read_only(&mut self, nodes: &[TensorId]) -> Vec<TensorIr> {
impl core::hash::Hash for InitOperationIr {
    fn hash<H: core::hash::Hasher>(&self, state: &mut H) {
impl InitOperationIr {
    fn nodes(&self) -> Vec<&TensorIr> {
impl TensorIr {
    fn mark_read_only(&mut self, nodes: &[TensorId], output: &mut Vec<TensorIr>) {
impl core::hash::Hash for RandomOpIr {
    fn hash<H: core::hash::Hasher>(&self, state: &mut H) {
impl<E> core::hash::Hash for ScalarOpIr<E> {
    fn hash<H: core::hash::Hasher>(&self, state: &mut H) {
impl<E> core::hash::Hash for MaskFillOpIr<E> {
    fn hash<H: core::hash::Hasher>(&self, state: &mut H) {
impl<E> core::hash::Hash for ClampOpIr<E> {
    fn hash<H: core::hash::Hasher>(&self, state: &mut H) {
impl<E> core::hash::Hash for NumericOperationIr<E> {
    fn hash<H: core::hash::Hasher>(&self, state: &mut H) {
/// The tensor unique identifier.
pub struct TensorId {
impl core::fmt::Display for TensorId {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
/// The status of the current tensor.
pub enum TensorStatus {
/// The tensor can be read, but not written.
/// The tensor can be mutated inplace.
/// No handle exists for that tensor.
/// A tensor definition represents a snapshot of a tensor when it was used.
///
/// # Example
///
/// A tensor that is used multiple times has its status updated for each operation.
///
///   1. Status::NotInit
///   2. Status::ReadOnly
///   3. Status::ReadOnly
///   4. Status::ReadWrite
pub struct TensorIr {
/// The [tensor id](TensorId).
/// The shape of the tensor.
/// The [status](TensorStatus) of the tensor when it was used.
/// The [type](DType) of the tensor.
impl TensorId {
/// Create a new tensor id.
    pub fn new(value: u64) -> Self {
fn main() {
// https://github.com/rust-ndarray/ndarray/issues/1197
pub(crate) static SEED: Mutex<Option<StdRng>> = Mutex::new(None);
/// The device type for the ndarray backend.
pub enum NdArrayDevice {
/// The CPU device.
impl DeviceOps for NdArrayDevice {
    fn id(&self) -> burn_tensor::backend::DeviceId {
impl Default for NdArrayDevice {
    fn default() -> Self {
/// Tensor backend that uses the [ndarray](ndarray) crate for executing tensor operations.
///
/// This backend is compatible with CPUs and can be compiled for almost any platform, including
/// `wasm`, `arm`, and `x86`.
pub struct NdArray<E = f32, I = i64, Q = i8> {
impl<E: FloatNdArrayElement, I: IntNdArrayElement, Q: QuantElement> Backend for NdArray<E, I, Q> {
    type Device = NdArrayDevice;
    type FloatTensorPrimitive = NdArrayTensorFloat;
    type FloatElem = E;
    type IntTensorPrimitive = NdArrayTensor<I>;
    type IntElem = I;
    type BoolTensorPrimitive = NdArrayTensor<bool>;
    type BoolElem = bool;
    type QuantizedTensorPrimitive = NdArrayQTensor<Q>;
    type QuantizedEncoding = Q;
    fn ad_enabled() -> bool {
    fn name(_device: &Self::Device) -> String {
    fn seed(seed: u64) {
impl<E: FloatNdArrayElement, I: IntNdArrayElement, Q: QuantElement> BackendIr for NdArray<E, I, Q> {
    type Handle = HandleKind<Self>;
    fn float_tensor(handle: TensorHandle<Self::Handle>) -> FloatTensor<Self> {
    fn int_tensor(handle: TensorHandle<Self::Handle>) -> IntTensor<Self> {
    fn bool_tensor(handle: TensorHandle<Self::Handle>) -> BoolTensor<Self> {
    fn quantized_tensor(handle: TensorHandle<Self::Handle>) -> QuantizedTensor<Self> {
    fn float_tensor_handle(tensor: FloatTensor<Self>) -> Self::Handle {
    fn int_tensor_handle(tensor: IntTensor<Self>) -> Self::Handle {
    fn bool_tensor_handle(tensor: BoolTensor<Self>) -> Self::Handle {
    fn quantized_tensor_handle(tensor: QuantizedTensor<Self>) -> Self::Handle {
/// A float element for ndarray backend.
pub trait FloatNdArrayElement: NdArrayElement + LinalgScalar + Signed
where
    Self: Sized,
{
/// An int element for ndarray backend.
pub trait IntNdArrayElement: NdArrayElement + Signed {
/// A general element for ndarray backend.
pub trait NdArrayElement:
    Element
    + ndarray::LinalgScalar
    + ndarray::ScalarOperand
    + ExpElement
    + num_traits::FromPrimitive
    + core::ops::AddAssign
    + core::cmp::PartialEq
    + core::cmp::PartialOrd<Self>
    + core::ops::Rem<Output = Self>
{
/// A element for ndarray backend that supports exp ops.
pub trait ExpElement {
/// Exponent
    fn exp_elem(self) -> Self;
/// Log
    fn log_elem(self) -> Self;
/// Log1p
    fn log1p_elem(self) -> Self;
/// Powf
    fn powf_elem(self, value: f32) -> Self;
/// Powi
    fn powi_elem(self, value: i32) -> Self;
/// Sqrt
    fn sqrt_elem(self) -> Self;
/// Abs
    fn abs_elem(self) -> Self;
/// Abs for int
    fn int_abs_elem(self) -> Self;
/// A quantized element for the ndarray backend.
pub trait QuantElement: NdArrayElement {
impl QuantElement for i8 {
impl FloatNdArrayElement for f64 {
impl FloatNdArrayElement for f32 {
impl IntNdArrayElement for i64 {
impl IntNdArrayElement for i32 {
        impl NdArrayElement for $ty {
        impl ExpElement for $ty {
            fn exp_elem(self) -> Self {
            fn log_elem(self) -> Self {
            fn log1p_elem(self) -> Self {
            fn powf_elem(self, value: f32) -> Self {
            fn powi_elem(self, value: i32) -> Self {
            fn sqrt_elem(self) -> Self {
            fn abs_elem(self) -> Self {
            fn int_abs_elem(self) -> Self {
        impl NdArrayElement for $ty {
        impl ExpElement for $ty {
            fn exp_elem(self) -> Self {
            fn log_elem(self) -> Self {
            fn log1p_elem(self) -> Self {
            fn powf_elem(self, value: f32) -> Self {
            fn powi_elem(self, value: i32) -> Self {
            fn sqrt_elem(self) -> Self {
            fn abs_elem(self) -> Self {
            fn int_abs_elem(self) -> Self {
//! Burn ndarray backend.
    type TestBackend = crate::NdArray<f32>;
    type TestTensor<const D: usize> = burn_tensor::Tensor<TestBackend, D>;
    type TestTensorInt<const D: usize> = burn_tensor::Tensor<TestBackend, D, burn_tensor::Int>;
    type TestTensorBool<const D: usize> = burn_tensor::Tensor<TestBackend, D, burn_tensor::Bool>;
impl<E: FloatNdArrayElement, I: IntNdArrayElement, Q: QuantElement> ActivationOps<Self>
    for NdArray<E, I, Q>
{
    fn relu(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
pub(crate) fn adaptive_avg_pool2d<E: FloatNdArrayElement>(
    x: NdArrayTensor<E>,
    output_size: [usize;
pub(crate) fn adaptive_avg_pool2d_backward<E: FloatNdArrayElement>(
    x: NdArrayTensor<E>,
    grad: NdArrayTensor<E>,
) -> NdArrayTensor<E> {
fn start_index(output_size_index: usize, output_size: usize, input_size: usize) -> usize {
fn end_index(output_size_index: usize, output_size: usize, input_size: usize) -> usize {
pub(crate) fn avg_pool2d<E: FloatNdArrayElement>(
    x: NdArrayTensor<E>,
    kernel_size: [usize;
pub(crate) fn avg_pool2d_backward<E: FloatNdArrayElement>(
    x: NdArrayTensor<E>,
    grad: NdArrayTensor<E>,
    kernel_size: [usize;
pub struct NdArrayOps<E> {
pub(crate) struct NdArrayMathOps<E> {
impl<E> NdArrayOps<E>
where
    E: Copy + Debug + burn_tensor::Element,
{
    pub fn into_data(tensor: NdArrayTensor<E>) -> TensorData {
    pub fn slice(tensor: NdArrayTensor<E>, ranges: &[Range<usize>]) -> NdArrayTensor<E> {
    pub fn slice_assign(
        tensor: NdArrayTensor<E>,
        ranges: &[Range<usize>],
        value: NdArrayTensor<E>,
    ) -> NdArrayTensor<E> {
    pub fn reshape(tensor: NdArrayTensor<E>, shape: Shape) -> NdArrayTensor<E> {
    pub(crate) fn concatenate(
        arrays: &[ndarray::ArrayView<E, IxDyn>],
        dim: usize,
    ) -> NdArrayTensor<E> {
// Transform column-major layout into row-major (standard) layout. (fix #1053)
    pub fn cat(tensors: Vec<NdArrayTensor<E>>, dim: usize) -> NdArrayTensor<E> {
    fn to_slice_args(ranges: &[Range<usize>], ndims: usize) -> Vec<SliceInfoElem> {
    pub fn swap_dims(tensor: NdArrayTensor<E>, dim1: usize, dim2: usize) -> NdArrayTensor<E> {
    pub fn permute(tensor: NdArrayTensor<E>, axes: &[usize]) -> NdArrayTensor<E> {
/// Broadcasts the tensor to the given shape
    pub(crate) fn expand(tensor: NdArrayTensor<E>, shape: Shape) -> NdArrayTensor<E> {
// need to convert view to owned array because NdArrayTensor expects owned array
// and try_into_owned_nocopy() panics for broadcasted arrays (zero strides)
    pub fn flip(tensor: NdArrayTensor<E>, axes: &[usize]) -> NdArrayTensor<E> {
impl<E> NdArrayMathOps<E>
where
    E: Copy + NdArrayElement,
{
    pub fn add(lhs: NdArrayTensor<E>, rhs: NdArrayTensor<E>) -> NdArrayTensor<E> {
    pub fn add_scalar(lhs: NdArrayTensor<E>, rhs: E) -> NdArrayTensor<E> {
    pub fn sub(lhs: NdArrayTensor<E>, rhs: NdArrayTensor<E>) -> NdArrayTensor<E> {
    pub fn sub_scalar(lhs: NdArrayTensor<E>, rhs: E) -> NdArrayTensor<E> {
    pub fn mul(lhs: NdArrayTensor<E>, rhs: NdArrayTensor<E>) -> NdArrayTensor<E> {
    pub fn mul_scalar(lhs: NdArrayTensor<E>, rhs: E) -> NdArrayTensor<E> {
    pub fn div(lhs: NdArrayTensor<E>, rhs: NdArrayTensor<E>) -> NdArrayTensor<E> {
    pub fn div_scalar(lhs: NdArrayTensor<E>, rhs: E) -> NdArrayTensor<E> {
    pub fn remainder(lhs: NdArrayTensor<E>, rhs: NdArrayTensor<E>) -> NdArrayTensor<E> {
    pub fn remainder_scalar(lhs: NdArrayTensor<E>, rhs: E) -> NdArrayTensor<E>
    where
        E: core::ops::Rem<Output = E>,
    {
    pub fn recip(tensor: NdArrayTensor<E>) -> NdArrayTensor<E> {
    pub fn mean(tensor: NdArrayTensor<E>) -> NdArrayTensor<E> {
    pub fn sum(tensor: NdArrayTensor<E>) -> NdArrayTensor<E> {
    pub fn prod(tensor: NdArrayTensor<E>) -> NdArrayTensor<E> {
    pub fn mean_dim(tensor: NdArrayTensor<E>, dim: usize) -> NdArrayTensor<E> {
    pub fn sum_dim(tensor: NdArrayTensor<E>, dim: usize) -> NdArrayTensor<E> {
    pub fn prod_dim(tensor: NdArrayTensor<E>, dim: usize) -> NdArrayTensor<E> {
    pub fn gather<I: NdArrayElement>(
        dim: usize,
        mut tensor: NdArrayTensor<E>,
        mut indices: NdArrayTensor<I>,
    ) -> NdArrayTensor<E> {
    pub fn scatter<I: NdArrayElement>(
        dim: usize,
        mut tensor: NdArrayTensor<E>,
        mut indices: NdArrayTensor<I>,
        mut value: NdArrayTensor<E>,
    ) -> NdArrayTensor<E> {
    pub fn mask_where(
        tensor: NdArrayTensor<E>,
        mask: NdArrayTensor<bool>,
        source: NdArrayTensor<E>,
    ) -> NdArrayTensor<E> {
    pub fn mask_fill(
        tensor: NdArrayTensor<E>,
        mask: NdArrayTensor<bool>,
        value: E,
    ) -> NdArrayTensor<E> {
    fn gather_batch_size(shape_tensor: &Shape, shape_indices: &Shape) -> usize {
    pub fn select<I: NdArrayElement>(
        tensor: NdArrayTensor<E>,
        dim: usize,
        indices: NdArrayTensor<I>,
    ) -> NdArrayTensor<E> {
    pub fn select_assign<I: NdArrayElement>(
        tensor: NdArrayTensor<E>,
        dim: usize,
        indices: NdArrayTensor<I>,
        value: NdArrayTensor<E>,
    ) -> NdArrayTensor<E> {
    pub fn argmax<I: NdArrayElement>(tensor: NdArrayTensor<E>, dim: usize) -> NdArrayTensor<I> {
    pub fn argmin<I: NdArrayElement>(tensor: NdArrayTensor<E>, dim: usize) -> NdArrayTensor<I> {
    pub fn clamp_min(tensor: NdArrayTensor<E>, min: E) -> NdArrayTensor<E> {
    pub fn clamp_max(tensor: NdArrayTensor<E>, max: E) -> NdArrayTensor<E> {
    pub fn clamp(tensor: NdArrayTensor<E>, min: E, max: E) -> NdArrayTensor<E> {
    pub(crate) fn elementwise_op<OtherE>(
        lhs: NdArrayTensor<E>,
        rhs: NdArrayTensor<OtherE>,
        var_name: impl FnMut(&E, &OtherE) -> E,
    ) -> NdArrayTensor<E> {
    pub(crate) fn elementwise_op_scalar(
        lhs: NdArrayTensor<E>,
        var_name: impl FnMut(E) -> E,
    ) -> NdArrayTensor<E> {
    pub(crate) fn sign_op(tensor: NdArrayTensor<E>) -> NdArrayTensor<E>
    where
        E: Signed,
    {
    pub(crate) fn abs(tensor: NdArrayTensor<E>) -> NdArrayTensor<E> {
    pub(crate) fn equal(lhs: NdArrayTensor<E>, rhs: NdArrayTensor<E>) -> NdArrayTensor<bool> {
    pub(crate) fn equal_elem(lhs: NdArrayTensor<E>, rhs: E) -> NdArrayTensor<bool> {
    pub(crate) fn greater(lhs: NdArrayTensor<E>, rhs: NdArrayTensor<E>) -> NdArrayTensor<bool> {
    pub(crate) fn greater_elem(lhs: NdArrayTensor<E>, rhs: E) -> NdArrayTensor<bool> {
    pub(crate) fn greater_equal(
        lhs: NdArrayTensor<E>,
        rhs: NdArrayTensor<E>,
    ) -> NdArrayTensor<bool> {
    pub(crate) fn greater_equal_elem(lhs: NdArrayTensor<E>, rhs: E) -> NdArrayTensor<bool> {
    pub(crate) fn lower_equal(lhs: NdArrayTensor<E>, rhs: NdArrayTensor<E>) -> NdArrayTensor<bool> {
    pub(crate) fn lower_equal_elem(lhs: NdArrayTensor<E>, rhs: E) -> NdArrayTensor<bool> {
    pub(crate) fn lower(lhs: NdArrayTensor<E>, rhs: NdArrayTensor<E>) -> NdArrayTensor<bool> {
    pub(crate) fn lower_elem(lhs: NdArrayTensor<E>, rhs: E) -> NdArrayTensor<bool> {
pub struct NdArrayBitOps<I: IntNdArrayElement>(PhantomData<I>);
impl<I: IntNdArrayElement> NdArrayBitOps<I> {
    pub(crate) fn bitand(lhs: NdArrayTensor<I>, rhs: NdArrayTensor<I>) -> NdArrayTensor<I> {
    pub(crate) fn bitand_scalar(lhs: NdArrayTensor<I>, rhs: I) -> NdArrayTensor<I> {
    pub(crate) fn bitor(lhs: NdArrayTensor<I>, rhs: NdArrayTensor<I>) -> NdArrayTensor<I> {
    pub(crate) fn bitor_scalar(lhs: NdArrayTensor<I>, rhs: I) -> NdArrayTensor<I> {
    pub(crate) fn bitxor(lhs: NdArrayTensor<I>, rhs: NdArrayTensor<I>) -> NdArrayTensor<I> {
    pub(crate) fn bitxor_scalar(lhs: NdArrayTensor<I>, rhs: I) -> NdArrayTensor<I> {
    pub(crate) fn bitnot(tensor: NdArrayTensor<I>) -> NdArrayTensor<I> {
pub struct NdArrayBoolOps;
// Rust booleans are either `00000000` or `00000001`, so bitwise and/or is fine, but bitwise not would
// produce invalid values.
impl NdArrayBoolOps {
    pub(crate) fn equal(lhs: NdArrayTensor<bool>, rhs: NdArrayTensor<bool>) -> NdArrayTensor<bool> {
    pub(crate) fn and(lhs: NdArrayTensor<bool>, rhs: NdArrayTensor<bool>) -> NdArrayTensor<bool> {
    pub(crate) fn or(lhs: NdArrayTensor<bool>, rhs: NdArrayTensor<bool>) -> NdArrayTensor<bool> {
enum CmpType {
fn arg<E: NdArrayElement, I: NdArrayElement>(
    tensor: NdArrayTensor<E>,
    dim: usize,
    cmp: CmpType,
) -> NdArrayTensor<I> {
// Find the min/max value in the array, and return its index.
    fn should_generate_row_major_layout_for_cat() {
// unsqueeze dim on the outermost axis
// make `ndarray` concatenates array on the outermost axis
// Language
// Current crate
// Workspace crates
impl<E: FloatNdArrayElement, I: IntNdArrayElement, Q: QuantElement> BoolTensorOps<Self>
    for NdArray<E, I, Q>
{
    fn bool_from_data(data: TensorData, _device: &NdArrayDevice) -> NdArrayTensor<bool> {
    async fn bool_into_data(tensor: NdArrayTensor<bool>) -> TensorData {
    fn bool_to_device(tensor: NdArrayTensor<bool>, _device: &NdArrayDevice) -> NdArrayTensor<bool> {
    fn bool_reshape(tensor: NdArrayTensor<bool>, shape: Shape) -> NdArrayTensor<bool> {
    fn bool_slice(tensor: NdArrayTensor<bool>, ranges: &[Range<usize>]) -> NdArrayTensor<bool> {
    fn bool_into_int(tensor: NdArrayTensor<bool>) -> NdArrayTensor<I> {
    fn bool_device(_tensor: &NdArrayTensor<bool>) -> <NdArray<E> as Backend>::Device {
    fn bool_empty(shape: Shape, _device: &<NdArray<E> as Backend>::Device) -> NdArrayTensor<bool> {
    fn bool_slice_assign(
        tensor: NdArrayTensor<bool>,
        ranges: &[Range<usize>],
        value: NdArrayTensor<bool>,
    ) -> NdArrayTensor<bool> {
    fn bool_cat(tensors: Vec<NdArrayTensor<bool>>, dim: usize) -> NdArrayTensor<bool> {
    fn bool_equal(lhs: NdArrayTensor<bool>, rhs: NdArrayTensor<bool>) -> NdArrayTensor<bool> {
    fn bool_not(tensor: NdArrayTensor<bool>) -> NdArrayTensor<bool> {
    fn bool_and(lhs: NdArrayTensor<bool>, rhs: NdArrayTensor<bool>) -> NdArrayTensor<bool> {
    fn bool_or(lhs: NdArrayTensor<bool>, rhs: NdArrayTensor<bool>) -> NdArrayTensor<bool> {
    fn bool_into_float(tensor: NdArrayTensor<bool>) -> FloatTensor<Self> {
    fn bool_swap_dims(
        tensor: NdArrayTensor<bool>,
        dim1: usize,
        dim2: usize,
    ) -> NdArrayTensor<bool> {
    fn bool_permute(tensor: NdArrayTensor<bool>, axes: &[usize]) -> NdArrayTensor<bool> {
    fn bool_expand(tensor: NdArrayTensor<bool>, shape: Shape) -> NdArrayTensor<bool> {
    fn bool_flip(tensor: NdArrayTensor<bool>, axes: &[usize]) -> NdArrayTensor<bool> {
fn conv2d_mad_inner<E: FloatNdArrayElement>(
    mut output: ArrayViewMut2<E>,
    x: ArrayView2<E>,
    k: E,
    k_xy: (usize, usize),
    out_xy: (usize, usize),
    stride: (usize, usize),
    dilation: (usize, usize),
) {
// Construct a sub-slice view of the input row.
// This is done upfront so that rustc does not have to emit bounds checks
// in the hot loop below.
// Ditto. Construct a sub-slice view of the output row, and explicitly specify
// the bounds upfront as 0..out_width so that rustc can make the assumption
// that all accesses are in-bounds in the below loop.
fn conv3d_mad_inner<E: FloatNdArrayElement>(
    mut output: ArrayViewMut3<E>,
    x: ArrayView3<E>,
    k: E,
    k_xyz: (usize, usize, usize),
    out_xyz: (usize, usize, usize),
    stride: (usize, usize, usize),
    dilation: (usize, usize, usize),
) {
// Construct a sub-slice view of the input row.
// This is done upfront so that rustc does not have to emit bounds checks
// in the hot loop below.
// Ditto. Construct a sub-slice view of the output row, and explicitly specify
// the bounds upfront as 0..out_width so that rustc can make the assumption
// that all accesses are in-bounds in the below loop.
pub(crate) fn conv2d<E: FloatNdArrayElement>(
    x: NdArrayTensor<E>,
    weight: NdArrayTensor<E>,
    bias: Option<NdArrayTensor<E>>,
    options: ConvOptions<2>,
) -> NdArrayTensor<E> {
// Convert inputs from dynamic indexes to static to improve perf.
// NOTE: This function call is duplicated twice so that the compiler can perform auto-vectorization
// in the case that the stride/dilation is 1.
// Get a mutable slice reference to the row we're looping over.
// We explicitly define the bounds to 0..out_width so that rustc can make
// the assumption that all accesses are in-bounds.
pub(crate) fn conv_transpose2d<E: FloatNdArrayElement>(
    x: NdArrayTensor<E>,
    weight: NdArrayTensor<E>,
    bias: Option<NdArrayTensor<E>>,
    options: ConvTransposeOptions<2>,
) -> NdArrayTensor<E> {
pub(crate) fn conv3d<E: FloatNdArrayElement>(
    x: NdArrayTensor<E>,
    weight: NdArrayTensor<E>,
    bias: Option<NdArrayTensor<E>>,
    options: ConvOptions<3>,
) -> NdArrayTensor<E> {
// Convert inputs from dynamic indexes to static to improve perf.
// NOTE: This function call is duplicated twice so that the compiler can perform auto-vectorization
// in the case that the stride/dilation is 1.
// Get a mutable iterator to the row we're looping over.
// We explicitly define the bounds to 0..out_width so that rustc can make
// the assumption that all accesses are in-bounds.
pub(crate) fn conv_transpose3d<E: FloatNdArrayElement>(
    x: NdArrayTensor<E>,
    weight: NdArrayTensor<E>,
    bias: Option<NdArrayTensor<E>>,
    options: ConvTransposeOptions<3>,
) -> NdArrayTensor<E> {
fn deform_im2col_kernel<F: FloatNdArrayElement>(
    out_y: usize,
    out_x: usize,
    input: ArrayView2<F>,
    offset: ArrayView3<F>,
    mask: Option<ArrayView2<F>>,
    mut columns: ArrayViewMut2<F>,
    args: DeformConvOptions<2>,
    (kernel_h, kernel_w): (usize, usize),
) {
// position shape: [in_channels, batch_size, out_h, out_w]
// columns shape: [[in_channels, kernel_h, kernel_w], [batch_size, out_h, out_w]]
fn bilinear_interpolate<F: FloatNdArrayElement>(
    input: ArrayView2<F>,
    height: usize,
    width: usize,
    y: F,
    x: F,
) -> F {
// To simplify code
pub(crate) fn deform_conv2d<F: FloatNdArrayElement>(
    input: NdArrayTensor<F>,
    offset: NdArrayTensor<F>,
    weight: NdArrayTensor<F>,
    mask: Option<NdArrayTensor<F>>,
    bias: Option<NdArrayTensor<F>>,
    args: DeformConvOptions<2>,
) -> NdArrayTensor<F> {
pub(crate) fn deform_im2col<F: FloatNdArrayElement>(
    input: ArrayView4<F>,
    offset: ArrayView4<F>,
    mask: Option<ArrayView6<F>>,
    args: DeformConvOptions<2>,
    out_dims: (usize, usize),
    kernel_dims: (usize, usize),
) -> Array2<F> {
// Columns is created here, so we know it's contiguous
    pub(crate) type DeformConv2dBackward<F> = (
        NdArrayTensor<F>,
        NdArrayTensor<F>,
        NdArrayTensor<F>,
        Option<NdArrayTensor<F>>,
        Option<NdArrayTensor<F>>,
    );
/// Calculate the [deformable 2D convolution](crate::ops::ModuleOps::deform_conv2d) backward pass using convolutions.
    pub(crate) fn deform_conv2d_backward<F: FloatNdArrayElement>(
        input: NdArrayTensor<F>,
        offset: NdArrayTensor<F>,
        weight: NdArrayTensor<F>,
        mask: Option<NdArrayTensor<F>>,
        bias: Option<NdArrayTensor<F>>,
        out_grad: NdArrayTensor<F>,
        args: DeformConvOptions<2>,
    ) -> DeformConv2dBackward<F> {
    fn compute_weight_grad<F: FloatNdArrayElement>(
        input: ArrayView4<F>,
        offset: ArrayView4<F>,
        mask: Option<ArrayView6<F>>,
        out_grad: ArrayView3<F>,
        options: DeformConvOptions<2>,
        kernel_dims: (usize, usize),
        out_dims: (usize, usize),
    ) -> NdArrayTensor<F> {
    type InputGradients<F> = (NdArrayTensor<F>, NdArrayTensor<F>, Option<NdArrayTensor<F>>);
    fn backward_gradient_inputs<F: FloatNdArrayElement>(
        image: ArrayView4<F>,
        weight: NdArrayTensor<F>,
        offset: ArrayView4<F>,
        mask: Option<ArrayView6<F>>,
        out_grad: ArrayView3<F>,
        args: &DeformConvOptions<2>,
        kernel_dims: (usize, usize),
    ) -> InputGradients<F> {
    fn compute_offset_and_mask_gradient<F: FloatNdArrayElement>(
        columns: ArrayView6<F>,
        image: ArrayView4<F>,
        offset: ArrayView4<F>,
        mask: Option<ArrayView6<F>>,
        args: &DeformConvOptions<2>,
        kernel_dims: (usize, usize),
    ) -> (NdArrayTensor<F>, Option<NdArrayTensor<F>>) {
    fn get_coordinate_weight<F: FloatNdArrayElement>(
        input: ArrayView2<F>,
        height: usize,
        width: usize,
        y: F,
        x: F,
        is_y_direction: bool,
    ) -> F {
    fn compute_input_grad<F: FloatNdArrayElement>(
        columns: ArrayView6<F>,
        offset: ArrayView4<F>,
        mask: Option<ArrayView6<F>>,
        args: &DeformConvOptions<2>,
        kernel_dims: (usize, usize),
        input_shape: (usize, usize, usize, usize),
    ) -> NdArrayTensor<F> {
// `for_each` expects a 2-tuple argument with `.into_par_iter()`, but 2 separate arguments otherwise
    fn deform_col2img_kernel(
        y: f32,
        x: f32,
        mask: Option<f32>,
        col: f32,
        grad_input: ArrayView2<AtomicF32>,
    ) {
// Language
// Current crate
// Workspace crates
impl<E: FloatNdArrayElement, I: IntNdArrayElement, Q: QuantElement> IntTensorOps<Self>
    for NdArray<E, I, Q>
{
    fn int_from_data(data: TensorData, _device: &NdArrayDevice) -> NdArrayTensor<I> {
    async fn int_into_data(tensor: NdArrayTensor<I>) -> TensorData {
    fn int_to_device(tensor: NdArrayTensor<I>, _device: &NdArrayDevice) -> NdArrayTensor<I> {
    fn int_reshape(tensor: NdArrayTensor<I>, shape: Shape) -> NdArrayTensor<I> {
    fn int_slice(tensor: NdArrayTensor<I>, ranges: &[Range<usize>]) -> NdArrayTensor<I> {
    fn int_device(_tensor: &NdArrayTensor<I>) -> <NdArray<E> as Backend>::Device {
    fn int_empty(shape: Shape, device: &<NdArray<E> as Backend>::Device) -> NdArrayTensor<I> {
    fn int_mask_where(
        tensor: NdArrayTensor<I>,
        mask: NdArrayTensor<bool>,
        source: NdArrayTensor<I>,
    ) -> NdArrayTensor<I> {
    fn int_mask_fill(
        tensor: NdArrayTensor<I>,
        mask: NdArrayTensor<bool>,
        value: I,
    ) -> NdArrayTensor<I> {
    fn int_slice_assign(
        tensor: NdArrayTensor<I>,
        ranges: &[Range<usize>],
        value: NdArrayTensor<I>,
    ) -> NdArrayTensor<I> {
    fn int_cat(tensors: Vec<NdArrayTensor<I>>, dim: usize) -> NdArrayTensor<I> {
    fn int_equal(lhs: NdArrayTensor<I>, rhs: NdArrayTensor<I>) -> NdArrayTensor<bool> {
    fn int_equal_elem(lhs: NdArrayTensor<I>, rhs: I) -> NdArrayTensor<bool> {
    fn int_greater(lhs: NdArrayTensor<I>, rhs: NdArrayTensor<I>) -> NdArrayTensor<bool> {
    fn int_greater_elem(lhs: NdArrayTensor<I>, rhs: I) -> NdArrayTensor<bool> {
    fn int_greater_equal(lhs: NdArrayTensor<I>, rhs: NdArrayTensor<I>) -> NdArrayTensor<bool> {
    fn int_greater_equal_elem(lhs: NdArrayTensor<I>, rhs: I) -> NdArrayTensor<bool> {
    fn int_lower(lhs: NdArrayTensor<I>, rhs: NdArrayTensor<I>) -> NdArrayTensor<bool> {
    fn int_lower_elem(lhs: NdArrayTensor<I>, rhs: I) -> NdArrayTensor<bool> {
    fn int_lower_equal(lhs: NdArrayTensor<I>, rhs: NdArrayTensor<I>) -> NdArrayTensor<bool> {
    fn int_lower_equal_elem(lhs: NdArrayTensor<I>, rhs: I) -> NdArrayTensor<bool> {
    fn int_add(lhs: NdArrayTensor<I>, rhs: NdArrayTensor<I>) -> NdArrayTensor<I> {
    fn int_add_scalar(lhs: NdArrayTensor<I>, rhs: I) -> NdArrayTensor<I> {
    fn int_sub(lhs: NdArrayTensor<I>, rhs: NdArrayTensor<I>) -> NdArrayTensor<I> {
    fn int_sub_scalar(lhs: NdArrayTensor<I>, rhs: I) -> NdArrayTensor<I> {
    fn int_mul(lhs: NdArrayTensor<I>, rhs: NdArrayTensor<I>) -> NdArrayTensor<I> {
    fn int_mul_scalar(lhs: NdArrayTensor<I>, rhs: I) -> NdArrayTensor<I> {
    fn int_div(lhs: NdArrayTensor<I>, rhs: NdArrayTensor<I>) -> NdArrayTensor<I> {
    fn int_div_scalar(lhs: NdArrayTensor<I>, rhs: I) -> NdArrayTensor<I> {
    fn int_remainder(lhs: NdArrayTensor<I>, rhs: NdArrayTensor<I>) -> NdArrayTensor<I> {
    fn int_remainder_scalar(lhs: NdArrayTensor<I>, rhs: I) -> NdArrayTensor<I> {
    fn int_neg(tensor: NdArrayTensor<I>) -> NdArrayTensor<I> {
    fn int_zeros(shape: Shape, device: &<NdArray<E> as Backend>::Device) -> NdArrayTensor<I> {
    fn int_ones(shape: Shape, device: &<NdArray<E> as Backend>::Device) -> NdArrayTensor<I> {
    fn int_full(
        shape: Shape,
        fill_value: I,
        device: &<NdArray<E> as Backend>::Device,
    ) -> NdArrayTensor<I> {
    fn int_sum(tensor: NdArrayTensor<I>) -> NdArrayTensor<I> {
    fn int_sum_dim(tensor: NdArrayTensor<I>, dim: usize) -> NdArrayTensor<I> {
    fn int_prod(tensor: NdArrayTensor<I>) -> NdArrayTensor<I> {
    fn int_prod_dim(tensor: NdArrayTensor<I>, dim: usize) -> NdArrayTensor<I> {
    fn int_mean(tensor: NdArrayTensor<I>) -> NdArrayTensor<I> {
    fn int_mean_dim(tensor: NdArrayTensor<I>, dim: usize) -> NdArrayTensor<I> {
    fn int_gather(
        dim: usize,
        tensor: NdArrayTensor<I>,
        indices: NdArrayTensor<I>,
    ) -> NdArrayTensor<I> {
    fn int_scatter(
        dim: usize,
        tensor: NdArrayTensor<I>,
        indices: NdArrayTensor<I>,
        value: NdArrayTensor<I>,
    ) -> NdArrayTensor<I> {
    fn int_select(
        tensor: NdArrayTensor<I>,
        dim: usize,
        indices: NdArrayTensor<I>,
    ) -> NdArrayTensor<I> {
    fn int_select_assign(
        tensor: NdArrayTensor<I>,
        dim: usize,
        indices: NdArrayTensor<I>,
        value: NdArrayTensor<I>,
    ) -> NdArrayTensor<I> {
    fn int_argmax(tensor: NdArrayTensor<I>, dim: usize) -> NdArrayTensor<I> {
    fn int_argmin(tensor: NdArrayTensor<I>, dim: usize) -> NdArrayTensor<I> {
    fn int_clamp_min(tensor: NdArrayTensor<I>, min: I) -> NdArrayTensor<I> {
    fn int_clamp_max(tensor: NdArrayTensor<I>, max: I) -> NdArrayTensor<I> {
    fn int_clamp(tensor: NdArrayTensor<I>, min: I, max: I) -> NdArrayTensor<I> {
    fn int_abs(tensor: NdArrayTensor<I>) -> NdArrayTensor<I> {
    fn int_into_float(tensor: NdArrayTensor<I>) -> FloatTensor<Self> {
    fn int_swap_dims(tensor: NdArrayTensor<I>, dim1: usize, dim2: usize) -> NdArrayTensor<I> {
    fn int_random(
        shape: Shape,
        distribution: Distribution,
        device: &NdArrayDevice,
    ) -> NdArrayTensor<I> {
// Assuming UniformInt is the integer variant
    fn int_powi(lhs: NdArrayTensor<I>, rhs: NdArrayTensor<I>) -> NdArrayTensor<I> {
    fn int_powf(lhs: NdArrayTensor<I>, rhs: FloatTensor<Self>) -> NdArrayTensor<I> {
    fn int_powf_scalar(lhs: NdArrayTensor<I>, rhs: f32) -> NdArrayTensor<I> {
    fn int_permute(tensor: NdArrayTensor<I>, axes: &[usize]) -> NdArrayTensor<I> {
    fn int_flip(tensor: NdArrayTensor<I>, axes: &[usize]) -> NdArrayTensor<I> {
    fn int_sign(tensor: NdArrayTensor<I>) -> NdArrayTensor<I> {
    fn int_expand(tensor: NdArrayTensor<I>, shape: Shape) -> NdArrayTensor<I> {
    fn bitwise_and(lhs: NdArrayTensor<I>, rhs: NdArrayTensor<I>) -> NdArrayTensor<I> {
    fn bitwise_and_scalar(lhs: NdArrayTensor<I>, rhs: I) -> NdArrayTensor<I> {
    fn bitwise_or(lhs: NdArrayTensor<I>, rhs: NdArrayTensor<I>) -> NdArrayTensor<I> {
    fn bitwise_or_scalar(lhs: NdArrayTensor<I>, rhs: I) -> NdArrayTensor<I> {
    fn bitwise_xor(lhs: NdArrayTensor<I>, rhs: NdArrayTensor<I>) -> NdArrayTensor<I> {
    fn bitwise_xor_scalar(lhs: NdArrayTensor<I>, rhs: I) -> NdArrayTensor<I> {
    fn bitwise_not(tensor: NdArrayTensor<I>) -> NdArrayTensor<I> {
    fn bitwise_left_shift(lhs: NdArrayTensor<I>, rhs: NdArrayTensor<I>) -> NdArrayTensor<I> {
    fn bitwise_left_shift_scalar(lhs: NdArrayTensor<I>, rhs: I) -> NdArrayTensor<I> {
    fn bitwise_right_shift(lhs: NdArrayTensor<I>, rhs: NdArrayTensor<I>) -> NdArrayTensor<I> {
    fn bitwise_right_shift_scalar(lhs: NdArrayTensor<I>, rhs: I) -> NdArrayTensor<I> {
pub(crate) fn nearest_interpolate<E: FloatNdArrayElement>(
    x: NdArrayTensor<E>,
    output_size: [usize;
pub(crate) fn nearest_interpolate_backward<E: FloatNdArrayElement>(
    x: NdArrayTensor<E>,
    grad: NdArrayTensor<E>,
    output_size: [usize;
fn start_index(output_size_index: usize, output_size: usize, input_size: usize) -> usize {
pub(crate) fn bilinear_interpolate<E: FloatNdArrayElement>(
    x: NdArrayTensor<E>,
    output_size: [usize;
// We convert everything to `f64` for calculations and then back to `E` at the end.
pub(crate) fn bicubic_interpolate<E: FloatNdArrayElement>(
    x: NdArrayTensor<E>,
    output_size: [usize;
    fn cubic_interp1d(x0: f64, x1: f64, x2: f64, x3: f64, t: f64) -> f64 {
        fn cubic_convolution1(x: f64, a: f64) -> f64 {
        fn cubic_convolution2(x: f64, a: f64) -> f64 {
pub(crate) fn mean_dim<E: NdArrayElement>(
    tensor: NdArrayTensor<E>,
    dim: usize,
) -> NdArrayTensor<E> {
pub(crate) fn sum_dim<E: NdArrayElement>(tensor: NdArrayTensor<E>, dim: usize) -> NdArrayTensor<E> {
pub(crate) fn prod_dim<E: NdArrayElement>(
    tensor: NdArrayTensor<E>,
    dim: usize,
) -> NdArrayTensor<E> {
pub(crate) fn matmul<E>(lhs: NdArrayTensor<E>, rhs: NdArrayTensor<E>) -> NdArrayTensor<E>
where
    E: FloatNdArrayElement,
{
// # of left rows
// # of left cols and right rows
// # of right cols
// size of matrix component of left array
// size of matrix component of right array
// size of matrix component of output array
// Here, we:
//   1. Un-flatten the output batch into a component-based batch index.
//   2. Use the strides for left and right batch indices to convert it to a flattened
//      batch for left and right.
struct Strides {
impl Strides {
    fn new(strides: Vec<usize>) -> Self {
    fn unflatten(&self, linear_index: usize) -> Vec<usize> {
    fn flatten(&self, index: &Vec<usize>) -> usize {
/// Compute the (broadcasted) output shape of matrix multiplication, along with strides for
/// the non-matrix dimensions of all arrays.
///
/// # Arguments
/// * `lsh`: Shape of the first (left-hand) matrix multiplication argument.
/// * `rsh`: Shape of the second (right-hand) matrix multiplication argument.
///
/// # Panics
/// * If `D` is not at least 2.
/// * If the matrix multiplication dimensions (last 2) are incompatible.
/// * If any other dimension is not the same for both tensors, or equal to 1. (Any dimension where
///   one dim is equal to 1 is broadcast.)
fn output_shape(lsh: &Shape, rsh: &Shape) -> (Shape, Strides, Strides, Strides) {
// Fetch matrix dimensions and check compatibility.
// Set matrix dimensions of the output shape.
// Set other array dimensions, broadcasting as necessary.
// Compute the strides inline.
// Compatible dimensions are:
//   1. Both dimensions are equal.
//   2. One of the dimensions is equal to 1.
// both dimensions are equal
// broadcast the left
// broadcast the right
    impl Strides {
        fn empty() -> Self {
    fn test_output_shape() {
// plain matrix multiply
// matrix multiply with one extra stack dimension
// rank 3, broadcast left
// rank 3, broadcast right
// rank 4, multi broadcast
// rank 5, multi-broadcast
    fn test_output_shape_too_small() {
    fn test_output_shape_bad_matrix_dims() {
    fn test_output_shape_non_broadcast() {
pub(crate) fn max_pool2d<E: FloatNdArrayElement>(
    x: NdArrayTensor<E>,
    kernel_size: [usize;
pub(crate) fn max_pool2d_with_indices<E: FloatNdArrayElement, I: IntNdArrayElement>(
    x: NdArrayTensor<E>,
    kernel_size: [usize;
pub(crate) fn max_pool2d_backward<E: FloatNdArrayElement, I: IntNdArrayElement>(
    x: NdArrayTensor<E>,
    _kernel_size: [usize;
// Module op with inputs (inp), optional (opt) and arguments (args).
                type $element = f32;
                type $element = f64;
impl<E: FloatNdArrayElement, I: IntNdArrayElement, Q: QuantElement> ModuleOps<Self>
    for NdArray<E, I, Q>
{
    fn conv2d(
        x: NdArrayTensorFloat,
        weight: NdArrayTensorFloat,
        bias: Option<NdArrayTensorFloat>,
        options: ConvOptions<2>,
    ) -> NdArrayTensorFloat {
    fn deform_conv2d(
        x: FloatTensor<Self>,
        offset: FloatTensor<Self>,
        weight: FloatTensor<Self>,
        mask: Option<FloatTensor<Self>>,
        bias: Option<FloatTensor<Self>>,
        options: DeformConvOptions<2>,
    ) -> FloatTensor<Self> {
    fn deform_conv2d_backward(
        x: FloatTensor<Self>,
        offset: FloatTensor<Self>,
        weight: FloatTensor<Self>,
        mask: Option<FloatTensor<Self>>,
        bias: Option<FloatTensor<Self>>,
        output_grad: FloatTensor<Self>,
        options: DeformConvOptions<2>,
    ) -> DeformConv2dBackward<Self> {
    fn conv_transpose2d(
        x: FloatTensor<Self>,
        weight: FloatTensor<Self>,
        bias: Option<FloatTensor<Self>>,
        options: ConvTransposeOptions<2>,
    ) -> FloatTensor<Self> {
    fn avg_pool2d(
        x: FloatTensor<Self>,
        kernel_size: [usize;
    fn avg_pool2d_backward(
        x: FloatTensor<Self>,
        grad: FloatTensor<Self>,
        kernel_size: [usize;
    fn max_pool2d(
        x: FloatTensor<Self>,
        kernel_size: [usize;
    fn max_pool2d_with_indices(
        x: FloatTensor<Self>,
        kernel_size: [usize;
    fn max_pool2d_with_indices_backward(
        x: FloatTensor<Self>,
        kernel_size: [usize;
    fn adaptive_avg_pool2d(x: FloatTensor<Self>, output_size: [usize;
    fn adaptive_avg_pool2d_backward(
        x: FloatTensor<Self>,
        grad: FloatTensor<Self>,
    ) -> FloatTensor<Self> {
    fn interpolate(
        x: FloatTensor<Self>,
        output_size: [usize;
    fn interpolate_backward(
        x: FloatTensor<Self>,
        grad: FloatTensor<Self>,
        output_size: [usize;
    fn conv3d(
        x: FloatTensor<Self>,
        weight: FloatTensor<Self>,
        bias: Option<FloatTensor<Self>>,
        options: ConvOptions<3>,
    ) -> FloatTensor<Self> {
    fn conv_transpose3d(
        x: FloatTensor<Self>,
        weight: FloatTensor<Self>,
        bias: Option<FloatTensor<Self>>,
        options: ConvTransposeOptions<3>,
    ) -> FloatTensor<Self> {
pub(crate) fn apply_padding_4d<E: Element>(
    x: NdArrayTensor<E>,
    padding: [usize;
pub(crate) fn apply_padding_5d<E: FloatNdArrayElement>(
    x: NdArrayTensor<E>,
    padding: [usize;
fn into_data<E: NdArrayElement>(tensor: NdArrayTensor<E>) -> TensorData {
fn into_data_f(tensor: NdArrayTensorFloat) -> TensorData {
impl<E: FloatNdArrayElement, I: IntNdArrayElement, Q: QuantElement> QTensorOps<Self>
    for NdArray<E, I, Q>
{
    fn q_from_data(data: TensorData, _device: &NdArrayDevice) -> QuantizedTensor<Self> {
// We should probably check that `Q` matches i8.. but it's the only valid type now
    fn quantize(
        tensor: FloatTensor<Self>,
        scheme: &QuantScheme,
        qparams: QuantizationParametersPrimitive<Self>,
    ) -> QuantizedTensor<Self> {
// Implement with ndarray instead of QuantizationStrategy?
    fn dequantize(tensor: QuantizedTensor<Self>) -> FloatTensor<Self> {
    fn q_device(_tensor: &QuantizedTensor<Self>) -> NdArrayDevice {
    fn q_to_device(
        tensor: QuantizedTensor<Self>,
        _device: &NdArrayDevice,
    ) -> QuantizedTensor<Self> {
    fn q_reshape(tensor: QuantizedTensor<Self>, shape: Shape) -> QuantizedTensor<Self> {
    async fn q_into_data(tensor: QuantizedTensor<Self>) -> TensorData {
    fn q_swap_dims(
        tensor: QuantizedTensor<Self>,
        dim1: usize,
        dim2: usize,
    ) -> QuantizedTensor<Self> {
    fn q_permute(tensor: QuantizedTensor<Self>, axes: &[usize]) -> QuantizedTensor<Self> {
    fn q_flip(tensor: QuantizedTensor<Self>, axes: &[usize]) -> QuantizedTensor<Self> {
    fn q_gather(
        dim: usize,
        tensor: QuantizedTensor<Self>,
        indices: IntTensor<Self>,
    ) -> QuantizedTensor<Self> {
    fn q_select(
        tensor: QuantizedTensor<Self>,
        dim: usize,
        indices: IntTensor<Self>,
    ) -> QuantizedTensor<Self> {
    fn q_slice(tensor: QuantizedTensor<Self>, ranges: &[Range<usize>]) -> QuantizedTensor<Self> {
    fn q_argmax(tensor: QuantizedTensor<Self>, dim: usize) -> IntTensor<Self> {
    fn q_argmin(tensor: QuantizedTensor<Self>, dim: usize) -> IntTensor<Self> {
    fn q_expand(tensor: QuantizedTensor<Self>, shape: Shape) -> QuantizedTensor<Self> {
fn is_accelerated<S: Simd, T: VAdd + VDiv>(_x: PhantomData<T>) -> bool {
pub(crate) fn try_avg_pool2d_simd<E: Element>(
    x: NdArrayTensor<E>,
    ksize: [usize;
// Strides must be unit, dilation isn't supported, rows must be contiguous
fn cast<T, E>(tensor: NdArrayTensor<T>) -> NdArrayTensor<E> {
// Until you can use associated constants as array size, we need to hardcode this.
// The most common config (x86-v3) has 16 registers, so use half of them for accumulators.
    const BLOCK_REGISTERS: usize = 8;
    pub(crate) fn avg_pool_nhwc<E: Element + VAdd + VDiv>(
        x: NdArrayTensor<E>,
        kernel_size: [usize;
// Floor division ensures `blocks * lanes * blocking factor` is always `<= out_channels`.
// An exclusive loop will always have `lanes * blocking factor` elements in bounds.
// Floor division means simd_end is always divisible by `lanes` and `<= out_channels`. An
// exclusive loop will always have `lanes` elements in bounds.
// SAFETY: Loop ranges are non-overlapping, so the unsafe shared reference is safe.
// SAFETY: See `loop_unblocked`
// SAFETY: Loop ranges are non-overlapping, so the unsafe shared reference is safe.
/// Execute the blocked (unrolled) portion of the pool.
    fn loop_blocked<'a, S: Simd, E: Element + VAdd + VDiv>(
        x: ArrayView3<'a, E>,
        mut out: ArrayViewMut3<'a, E>,
        kernel_size: [usize;
// If pixels are more than `padding` from the edges, the in pixel cannot be out of bounds
// SAFETY:
// Load a full vector from x[N * lanes]. This is bounds checked by the
// slice above.
// SAFETY:
// Store a full vector to out[N * lanes]. This is bounds checked by the
// slice above.
// Border pixels need bounds checks
// SAFETY:
// Load a full vector from x[N * lanes]. This is bounds checked by the
// slice above.
// SAFETY:
// Store a full vector to out[N * lanes]. This is bounds checked by the
// slice above.
/// Execute the unblocked (not unrolled) portion of the pool.
///
/// SAFETY: Safe as long as `ch + simd_lanes <= out_channels`.
    unsafe fn loop_unblocked<'a, S: Simd, E: Element + VAdd + VDiv>(
        x: ArrayView3<'a, E>,
        mut out: ArrayViewMut3<'a, E>,
        kernel_size: [usize;
// If pixels are not within padding range, bounds checks are always true
// Load a full vector from `x`. In bounds as long as `out_channels >= ch + lanes`
// Store a full vector to `out`. In bounds as long as `out_channels >= ch + lanes`.
// Border pixels need bounds checks
// Load a full vector from `x`. In bounds as long as `out_channels >= ch + lanes`
// Store a full vector to `out`. In bounds as long as `out_channels >= ch + lanes`.
/// Execute scalar portion of the pooling
    fn loop_scalar<E: Element + VAdd + VDiv>(
        x: ArrayView3<'_, E>,
        mut out: ArrayViewMut3<'_, E>,
        kernel_size: [usize;
// If pixels are not within padding range, bounds checks are always true
// Border pixels need bounds checks
/// Whether SIMD instructions are worth using
pub fn should_use_simd(len: usize) -> bool {
/// Whether SIMD instructions are worth using
pub fn should_use_simd(_len: usize) -> bool {
pub(crate) fn lanes<E: Scalar>() -> usize {
    struct lanes<__T0>(__T0);
    impl<E: Scalar> ::macerator::WithSimd for lanes<PhantomData<E>> {
        type Output = usize;
        fn with_simd<__S: ::macerator::Simd>(self) -> <Self as ::macerator::WithSimd>::Output {
fn lanes_simd<S: Simd, E: Scalar>(_ty: PhantomData<E>) -> usize {
pub trait MinMax {
    fn min(self, other: Self) -> Self;
    fn max(self, other: Self) -> Self;
        impl MinMax for $ty {
            fn min(self, other: Self) -> Self {
            fn max(self, other: Self) -> Self {
impl MinMax for f32 {
    fn min(self, other: Self) -> Self {
    fn max(self, other: Self) -> Self {
impl MinMax for f64 {
    fn min(self, other: Self) -> Self {
    fn max(self, other: Self) -> Self {
pub trait SimdBinop<T: Scalar, Out: Scalar> {
    fn apply_vec<S: Simd>(lhs: Vector<S, T>, rhs: Vector<S, T>) -> Vector<S, Out>;
    fn apply(lhs: T, rhs: T) -> Out;
    fn is_accelerated<S: Simd>() -> bool;
impl<T: VAdd> SimdBinop<T, T> for VecAdd {
    fn apply_vec<S: Simd>(lhs: Vector<S, T>, rhs: Vector<S, T>) -> Vector<S, T> {
    fn apply(lhs: T, rhs: T) -> T {
    fn is_accelerated<S: Simd>() -> bool {
impl<T: VDiv> SimdBinop<T, T> for VecDiv {
    fn apply_vec<S: Simd>(lhs: Vector<S, T>, rhs: Vector<S, T>) -> Vector<S, T> {
    fn apply(lhs: T, rhs: T) -> T {
    fn is_accelerated<S: Simd>() -> bool {
impl<T: VMul> SimdBinop<T, T> for VecMul {
    fn apply_vec<S: Simd>(lhs: Vector<S, T>, rhs: Vector<S, T>) -> Vector<S, T> {
    fn apply(lhs: T, rhs: T) -> T {
    fn is_accelerated<S: Simd>() -> bool {
impl<T: VSub> SimdBinop<T, T> for VecSub {
    fn apply_vec<S: Simd>(lhs: Vector<S, T>, rhs: Vector<S, T>) -> Vector<S, T> {
    fn apply(lhs: T, rhs: T) -> T {
    fn is_accelerated<S: Simd>() -> bool {
impl<T: VOrd + MinMax> SimdBinop<T, T> for VecMin {
    fn apply_vec<S: Simd>(lhs: Vector<S, T>, rhs: Vector<S, T>) -> Vector<S, T> {
    fn apply(lhs: T, rhs: T) -> T {
    fn is_accelerated<S: Simd>() -> bool {
impl<T: VOrd + MinMax> SimdBinop<T, T> for VecMax {
    fn apply_vec<S: Simd>(lhs: Vector<S, T>, rhs: Vector<S, T>) -> Vector<S, T> {
    fn apply(lhs: T, rhs: T) -> T {
    fn is_accelerated<S: Simd>() -> bool {
impl<T: VBitAnd> SimdBinop<T, T> for VecBitAnd {
    fn apply_vec<S: Simd>(lhs: Vector<S, T>, rhs: Vector<S, T>) -> Vector<S, T> {
    fn apply(lhs: T, rhs: T) -> T {
    fn is_accelerated<S: Simd>() -> bool {
impl<T: VBitOr> SimdBinop<T, T> for VecBitOr {
    fn apply_vec<S: Simd>(lhs: Vector<S, T>, rhs: Vector<S, T>) -> Vector<S, T> {
    fn apply(lhs: T, rhs: T) -> T {
    fn is_accelerated<S: Simd>() -> bool {
impl<T: VBitXor> SimdBinop<T, T> for VecBitXor {
    fn apply_vec<S: Simd>(lhs: Vector<S, T>, rhs: Vector<S, T>) -> Vector<S, T> {
    fn apply(lhs: T, rhs: T) -> T {
    fn is_accelerated<S: Simd>() -> bool {
fn is_accelerated<S: Simd, T: Scalar, Out: Scalar, Op: SimdBinop<T, Out>>(
    _x: PhantomData<(T, Out, Op)>,
) -> bool {
pub fn try_binary_simd<
    E: Element,
    EOut: Element,
    T: NdArrayElement + Scalar,
    Out: NdArrayElement + Scalar,
    Op: SimdBinop<T, Out>,
>(
    lhs: NdArrayTensor<E>,
    rhs: NdArrayTensor<E>,
) -> Result<NdArrayTensor<EOut>, (NdArrayTensor<E>, NdArrayTensor<E>)> {
// Used to assert traits based on the dynamic `DType`.
// Used to assert traits based on the dynamic `DType`.
fn binary_simd_same<
    T: NdArrayElement + Scalar,
    Out: NdArrayElement + Scalar,
    Op: SimdBinop<T, Out>,
>(
    lhs: NdArrayTensor<T>,
    rhs: NdArrayTensor<T>,
) -> NdArrayTensor<Out> {
fn binary<
    'a,
    S: Simd,
    T: NdArrayElement + Scalar,
    Out: NdArrayElement + Scalar,
    Op: SimdBinop<T, Out>,
>(
    lhs: &'a [T],
    rhs: &'a [T],
    out: &'a mut [Out],
    _op: PhantomData<Op>,
) where
    'a: 'a,
{
// Load one full vector from `lhs`.
// SAFETY: Guaranteed to be in bounds because `len == 8 * lanes`
// Load one full vector from `rhs`.
// SAFETY: Guaranteed to be in bounds because `len == 8 * lanes`
// Store one full vector to `out`.
// SAFETY: Guaranteed to be in bounds because `len == 8 * lanes`
// Load one full vector from `lhs`.
// SAFETY: Guaranteed to be in bounds because `len == lanes`
// Load one full vector from `rhs`.
// SAFETY: Guaranteed to be in bounds because `len == lanes`
// Store one full vector to `out`.
// SAFETY: Guaranteed to be in bounds because `len == lanes`
/// Unsafely alias a slice to use as an inline argument
fn unsafe_alias_slice_mut<'a, T>(slice: &mut [T]) -> &'a mut [T] {
pub trait ScalarSimdBinop<T: Scalar, Out: Scalar> {
    type Rhs: Copy;
    type RhsVec<S: Simd>: Copy;
    fn splat<S: Simd>(rhs: Self::Rhs) -> Self::RhsVec<S>;
    fn apply_vec<S: Simd>(lhs: Vector<S, T>, rhs: Self::RhsVec<S>) -> Vector<S, Out>;
    fn apply(lhs: T, rhs: Self::Rhs) -> Out;
    fn is_accelerated<S: Simd>() -> bool;
pub struct VecAdd;
pub struct VecDiv;
pub struct VecMul;
pub struct VecSub;
pub struct VecMin;
pub struct VecMax;
pub struct VecClamp;
pub struct VecBitAnd;
pub struct VecBitOr;
pub struct VecBitXor;
impl<T: VAdd> ScalarSimdBinop<T, T> for VecAdd {
    type Rhs = T;
    type RhsVec<S: Simd> = Vector<S, T>;
    fn splat<S: Simd>(rhs: Self::Rhs) -> Self::RhsVec<S> {
    fn apply_vec<S: Simd>(lhs: Vector<S, T>, rhs: Self::RhsVec<S>) -> Vector<S, T> {
    fn apply(lhs: T, rhs: T) -> T {
    fn is_accelerated<S: Simd>() -> bool {
impl<T: VDiv> ScalarSimdBinop<T, T> for VecDiv {
    type Rhs = T;
    type RhsVec<S: Simd> = Vector<S, T>;
    fn splat<S: Simd>(rhs: Self::Rhs) -> Self::RhsVec<S> {
    fn apply_vec<S: Simd>(lhs: Vector<S, T>, rhs: Self::RhsVec<S>) -> Vector<S, T> {
    fn apply(lhs: T, rhs: T) -> T {
    fn is_accelerated<S: Simd>() -> bool {
impl<T: VMul> ScalarSimdBinop<T, T> for VecMul {
    type Rhs = T;
    type RhsVec<S: Simd> = Vector<S, T>;
    fn splat<S: Simd>(rhs: Self::Rhs) -> Self::RhsVec<S> {
    fn apply_vec<S: Simd>(lhs: Vector<S, T>, rhs: Self::RhsVec<S>) -> Vector<S, T> {
    fn apply(lhs: T, rhs: T) -> T {
    fn is_accelerated<S: Simd>() -> bool {
impl<T: VSub> ScalarSimdBinop<T, T> for VecSub {
    type Rhs = T;
    type RhsVec<S: Simd> = Vector<S, T>;
    fn splat<S: Simd>(rhs: Self::Rhs) -> Self::RhsVec<S> {
    fn apply_vec<S: Simd>(lhs: Vector<S, T>, rhs: Self::RhsVec<S>) -> Vector<S, T> {
    fn apply(lhs: T, rhs: T) -> T {
    fn is_accelerated<S: Simd>() -> bool {
impl<T: VOrd + MinMax> ScalarSimdBinop<T, T> for VecMin {
    type Rhs = T;
    type RhsVec<S: Simd> = Vector<S, T>;
    fn splat<S: Simd>(rhs: Self::Rhs) -> Self::RhsVec<S> {
    fn apply_vec<S: Simd>(lhs: Vector<S, T>, rhs: Self::RhsVec<S>) -> Vector<S, T> {
    fn apply(lhs: T, rhs: T) -> T {
    fn is_accelerated<S: Simd>() -> bool {
impl<T: VOrd + MinMax> ScalarSimdBinop<T, T> for VecMax {
    type Rhs = T;
    type RhsVec<S: Simd> = Vector<S, T>;
    fn splat<S: Simd>(rhs: Self::Rhs) -> Self::RhsVec<S> {
    fn apply_vec<S: Simd>(lhs: Vector<S, T>, rhs: Self::RhsVec<S>) -> Vector<S, T> {
    fn apply(lhs: T, rhs: T) -> T {
    fn is_accelerated<S: Simd>() -> bool {
impl<T: VOrd + MinMax> ScalarSimdBinop<T, T> for VecClamp {
    type Rhs = (T, T);
    type RhsVec<S: Simd> = (Vector<S, T>, Vector<S, T>);
    fn splat<S: Simd>((min, max): Self::Rhs) -> Self::RhsVec<S> {
    fn apply_vec<S: Simd>(lhs: Vector<S, T>, (min, max): Self::RhsVec<S>) -> Vector<S, T> {
    fn apply(lhs: T, (min, max): Self::Rhs) -> T {
    fn is_accelerated<S: Simd>() -> bool {
impl<T: VBitAnd> ScalarSimdBinop<T, T> for VecBitAnd {
    type Rhs = T;
    type RhsVec<S: Simd> = Vector<S, T>;
    fn splat<S: Simd>(rhs: Self::Rhs) -> Self::RhsVec<S> {
    fn apply_vec<S: Simd>(lhs: Vector<S, T>, rhs: Self::RhsVec<S>) -> Vector<S, T> {
    fn apply(lhs: T, rhs: Self::Rhs) -> T {
    fn is_accelerated<S: Simd>() -> bool {
impl<T: VBitOr> ScalarSimdBinop<T, T> for VecBitOr {
    type Rhs = T;
    type RhsVec<S: Simd> = Vector<S, T>;
    fn splat<S: Simd>(rhs: Self::Rhs) -> Self::RhsVec<S> {
    fn apply_vec<S: Simd>(lhs: Vector<S, T>, rhs: Self::RhsVec<S>) -> Vector<S, T> {
    fn apply(lhs: T, rhs: Self::Rhs) -> T {
    fn is_accelerated<S: Simd>() -> bool {
impl<T: VBitXor> ScalarSimdBinop<T, T> for VecBitXor {
    type Rhs = T;
    type RhsVec<S: Simd> = Vector<S, T>;
    fn splat<S: Simd>(rhs: Self::Rhs) -> Self::RhsVec<S> {
    fn apply_vec<S: Simd>(lhs: Vector<S, T>, rhs: Self::RhsVec<S>) -> Vector<S, T> {
    fn apply(lhs: T, rhs: Self::Rhs) -> T {
    fn is_accelerated<S: Simd>() -> bool {
fn is_accelerated<S: Simd, T: Scalar, Out: Scalar, Op: ScalarSimdBinop<T, Out>>(
    _x: PhantomData<(T, Out, Op)>,
) -> bool {
pub fn try_binary_scalar_simd<
    E: NdArrayElement,
    EOut: NdArrayElement,
    T: NdArrayElement + Scalar,
    Out: NdArrayElement + Scalar,
    Op: ScalarSimdBinop<T, Out>,
>(
    input: NdArrayTensor<E>,
    elem: Op::Rhs,
) -> Result<NdArrayTensor<EOut>, NdArrayTensor<E>> {
// Used to assert traits based on the dynamic `DType`.
// Used to assert traits based on the dynamic `DType`.
/// Execute operation in place on an owned tensor
/// SAFETY:
/// Must ensure `size_of::<T> == size_of::<Out>` and `align_of::<T> >= align_of::<Out>`.
unsafe fn binary_scalar_simd_inplace<
    T: NdArrayElement + Scalar,
    Out: NdArrayElement + Scalar,
    Op: ScalarSimdBinop<T, Out>,
>(
    input: NdArrayTensor<T>,
    elem: Op::Rhs,
) -> NdArrayTensor<Out> {
// Buffer has the same elem size and is filled with the operation output, so this is safe
/// Create a new copy of the tensor as the output
fn binary_scalar_simd_owned<
    T: NdArrayElement + Scalar,
    Out: NdArrayElement + Scalar,
    Op: ScalarSimdBinop<T, Out>,
>(
    input: NdArrayTensor<T>,
    elem: Op::Rhs,
) -> NdArrayTensor<Out> {
fn binary_scalar_slice<
    'a,
    S: Simd,
    T: NdArrayElement + Scalar,
    Out: NdArrayElement + Scalar,
    Op: ScalarSimdBinop<T, Out>,
>(
    input: &'a [T],
    out: &'a mut [Out],
    rhs: Op::Rhs,
    _op: PhantomData<Op>,
) where
    'a: 'a,
{
// Load one full vector from `input`.
// SAFETY: Guaranteed to be in bounds because `len == 8 * lanes`
// Store one full vector to `out`.
// SAFETY: Guaranteed to be in bounds because `len == 8 * lanes`
// Load one full vector from `input`.
// SAFETY: Guaranteed to be in bounds because `len == lanes`
// Store one full vector to `out`.
// SAFETY: Guaranteed to be in bounds because `len == lanes`
/// Execute operation in line.
/// SAFETY:
/// Must ensure `size_of::<T> == size_of::<Out>` and `align_of::<T> >= align_of::<Out>`.
unsafe fn binary_scalar_slice_inplace<
    'a,
    S: Simd,
    T: NdArrayElement + Scalar,
    Out: NdArrayElement + Scalar,
    Op: ScalarSimdBinop<T, Out>,
>(
    buf: &'a mut [T],
    rhs: Op::Rhs,
    _op: PhantomData<(Out, Op)>,
) where
    'a: 'a,
{
// Load a full vector from the aligned portion of the buffer.
// SAFETY: `align_to_mut` guarantees we're aligned to `T::Vector`'s size, and there is
// always a full vector in bounds.
            let s~N = unsafe {
// Store a full vector at the same position as the input. Cast is safe because `Out` is
// size and align compatible
// Load a full vector from the aligned portion of the buffer.
// SAFETY: `align_to_mut` guarantees we're aligned to `T::Vector`'s size, and there is
// always a full vector in bounds.
        let s0 = unsafe {
// Store a full vector at the same position as the input. Cast is safe because `Out` is
// size and align compatible
pub trait SimdCmpOp<T: Scalar> {
    fn apply_vec<S: Simd>(lhs: Vector<S, T>, rhs: Vector<S, T>) -> T::Mask<S>;
    fn apply(lhs: T, rhs: T) -> bool;
    fn is_accelerated<S: Simd>() -> bool;
pub struct VecEquals;
impl<T: VEq> SimdCmpOp<T> for VecEquals {
    fn apply_vec<S: Simd>(lhs: Vector<S, T>, rhs: Vector<S, T>) -> <T as Scalar>::Mask<S> {
    fn apply(lhs: T, rhs: T) -> bool {
    fn is_accelerated<S: Simd>() -> bool {
pub struct VecGreater;
impl<T: VOrd + PartialOrd> SimdCmpOp<T> for VecGreater {
    fn apply_vec<S: Simd>(lhs: Vector<S, T>, rhs: Vector<S, T>) -> <T as Scalar>::Mask<S> {
    fn apply(lhs: T, rhs: T) -> bool {
    fn is_accelerated<S: Simd>() -> bool {
pub struct VecGreaterEq;
impl<T: VOrd + PartialOrd> SimdCmpOp<T> for VecGreaterEq {
    fn apply_vec<S: Simd>(lhs: Vector<S, T>, rhs: Vector<S, T>) -> <T as Scalar>::Mask<S> {
    fn apply(lhs: T, rhs: T) -> bool {
    fn is_accelerated<S: Simd>() -> bool {
pub struct VecLowerEq;
impl<T: VOrd + PartialOrd> SimdCmpOp<T> for VecLowerEq {
    fn apply_vec<S: Simd>(lhs: Vector<S, T>, rhs: Vector<S, T>) -> <T as Scalar>::Mask<S> {
    fn apply(lhs: T, rhs: T) -> bool {
    fn is_accelerated<S: Simd>() -> bool {
pub struct VecLower;
impl<T: VOrd + PartialOrd> SimdCmpOp<T> for VecLower {
    fn apply_vec<S: Simd>(lhs: Vector<S, T>, rhs: Vector<S, T>) -> <T as Scalar>::Mask<S> {
    fn apply(lhs: T, rhs: T) -> bool {
    fn is_accelerated<S: Simd>() -> bool {
fn is_accelerated<S: Simd, T: Scalar, Op: SimdCmpOp<T>>(_x: PhantomData<(T, Op)>) -> bool {
pub fn try_cmp_simd<E: Element, T: NdArrayElement + Scalar, Op: SimdCmpOp<T>>(
    lhs: NdArrayTensor<E>,
    rhs: NdArrayTensor<E>,
) -> Result<NdArrayTensor<bool>, (NdArrayTensor<E>, NdArrayTensor<E>)> {
// Used to assert traits based on the dynamic `DType`.
fn cmp_simd_same<T: NdArrayElement + Scalar, Op: SimdCmpOp<T>>(
    lhs: NdArrayTensor<T>,
    rhs: NdArrayTensor<T>,
) -> NdArrayTensor<bool> {
fn cmp<'a, S: Simd, T: NdArrayElement + Scalar, Op: SimdCmpOp<T>>(
    lhs: &'a [T],
    rhs: &'a [T],
    out: &'a mut [bool],
    _op: PhantomData<Op>,
) where
    'a: 'a,
{
// Load one full vector from `lhs`.
// SAFETY: Guaranteed to be in bounds because `len == 8 * lanes`
// Load one full vector from `rhs`.
// SAFETY: Guaranteed to be in bounds because `len == 8 * lanes`
// Store one full vector to `out`.
// SAFETY: Guaranteed to be in bounds because `len == 8 * lanes`
// Load one full vector from `lhs`.
// SAFETY: Guaranteed to be in bounds because `len == lanes`
// Load one full vector from `rhs`.
// SAFETY: Guaranteed to be in bounds because `len == lanes`
// Store one full vector to `out`.
// SAFETY: Guaranteed to be in bounds because `len == lanes`
/// Unsafely alias a slice to use as an inline argument
fn unsafe_alias_slice_mut<'a, T>(slice: &mut [T]) -> &'a mut [T] {
    pub fn try_cmp_scalar_simd<E: NdArrayElement, T: NdArrayElement + Scalar, Op: SimdCmpOp<T>>(
        input: NdArrayTensor<E>,
        elem: T,
    ) -> Result<NdArrayTensor<bool>, NdArrayTensor<E>> {
// Used to assert traits based on the dynamic `DType`.
/// Execute operation in place on an owned tensor
/// SAFETY:
/// Must ensure `size_of::<T> == size_of::<Out>` and `align_of::<T> >= align_of::<Out>`.
    unsafe fn cmp_scalar_simd_inplace<T: NdArrayElement + Scalar, Op: SimdCmpOp<T>>(
        input: NdArrayTensor<T>,
        elem: T,
    ) -> NdArrayTensor<bool> {
// Buffer has the same elem size and is filled with the operation output, so this is safe
/// Create a new copy of the tensor as the output
    fn cmp_scalar_simd_owned<T: NdArrayElement + Scalar, Op: SimdCmpOp<T>>(
        input: NdArrayTensor<T>,
        elem: T,
    ) -> NdArrayTensor<bool> {
    fn cmp_scalar_slice<'a, S: Simd, T: NdArrayElement + Scalar, Op: SimdCmpOp<T>>(
        input: &'a [T],
        out: &'a mut [bool],
        rhs: T,
        _op: PhantomData<Op>,
    ) where
        'a: 'a,
    {
// Load one full vector from `input`.
// SAFETY: Guaranteed to be in bounds because `len == 8 * lanes`
// Store one full vector to `out`.
// SAFETY: Guaranteed to be in bounds because `len == 8 * lanes`
// Load one full vector from `input`.
// SAFETY: Guaranteed to be in bounds because `len == lanes`
// Store one full vector to `out`.
// SAFETY: Guaranteed to be in bounds because `len == lanes`
/// Execute operation in line.
/// SAFETY:
/// Must ensure `size_of::<T> == size_of::<Out>` and `align_of::<T> >= align_of::<Out>`.
    unsafe fn cmp_scalar_slice_inplace<'a, S: Simd, T: NdArrayElement + Scalar, Op: SimdCmpOp<T>>(
        buf: &'a mut [T],
        rhs: T,
        _op: PhantomData<Op>,
    ) where
        'a: 'a,
    {
// Load a full vector from the aligned portion of the buffer.
// SAFETY: `align_to_mut` guarantees we're aligned to `T::Vector`'s size, and there is
// always a full vector in bounds.
                let s~N = unsafe {
// Store a full vector at the same position as the input. Cast is safe because `Out` is
// size and align compatible
// Load a full vector from the aligned portion of the buffer.
// SAFETY: `align_to_mut` guarantees we're aligned to `T::Vector`'s size, and there is
// always a full vector in bounds.
            let s0 = unsafe {
// Store a full vector at the same position as the input. Cast is safe because `Out` is
// size and align compatible
type Args<E> = (NdArrayTensor<E>, NdArrayTensor<E>, Option<NdArrayTensor<E>>);
pub fn try_conv2d_simd<E: FloatNdArrayElement>(
    x: NdArrayTensor<E>,
    weight: NdArrayTensor<E>,
    bias: Option<NdArrayTensor<E>>,
    options: ConvOptions<2>,
) -> Result<NdArrayTensor<E>, Args<E>> {
fn cast<T, E>(tensor: NdArrayTensor<T>) -> NdArrayTensor<E> {
/// Out-channel last SIMD accelerated direct convolution. Loop order and register blocking based on
/// E. Georganas, S. Avancha, K. Banerjee, D. Kalamkar, G. Henry, H. Pabst, A. Heinecke (2018).
/// Anatomy Of High-Performance Deep Learning Convolutions On SIMD Architectures.
/// SC '18, Article 6, pp. 1-12. arXiv:1808.05567. <https://arxiv.org/abs/1808.05567>.
fn conv2d<E: VMulAdd + Element, T: Element>(
    x: NdArrayTensor<T>,
    weight: NdArrayTensor<T>,
    bias: Option<NdArrayTensor<T>>,
    options: ConvOptions<2>,
    _ty: PhantomData<E>,
) -> Result<NdArrayTensor<T>, Args<T>> {
    fn precheck<S: Simd, E: VMulAdd>(_ty: PhantomData<E>) -> (usize, bool) {
// floor division means `(oc_blocks - 1) * lanes` can never be greater than `out_channels - lanes`.
// SAFETY: Slices are guaranteed to be non-overlapping, so having an unsafe shared reference
// is safe. `oc_blocks * lanes` must be `<= out_channels` to satisfy safety of inner function.
/// Size of register blocks, we need to hardcode this because Rust and the `seq` macro don't support
/// using associated constants as constant parameters. 8 works for all semi-modern CPUs but might
/// not be perfectly optimized for AVX-512 capable CPUs (which probably should use 16).
/// This should always be conservative, since oversizing it will cause register spills and that's
/// **much** worse than the performance lost with lower values.
const REGISTER_BLOCK: usize = 8;
/// Run a loop of conv2d.
/// # SAFETY
/// See `conv2d_inner_nopad`, `conv2d_inner_nopad_nostride`, `conv2d_remainder`.
/// Required preconditions: `ob * simd_lanes` must be `<= out_channels - simd_lanes`, `weights` and
/// `out` must have unit stride for the out channels.
unsafe fn conv2d_launch<
    'a,
    S: Simd,
    E: VMulAdd,
    const PAD: bool,
    const STRIDE: bool,
    const GROUPS: bool,
>(
    x: ArrayView3<'a, E>,
    weights: ArrayView4<'a, E>,
    bias: &'a Option<ArcArray1<E>>,
    out: &'a mut ArrayViewMut3<'a, E>,
    options: &'a ConvOptions<2>,
    ob: usize,
) where
    'a: 'a,
{
// Trick compiler into inlining 0 to padding
/// Execute the non-unrolled and/or padded portion of the convolution. This has more checks and is
/// much slower, so we want to minimize the amount of pixels that need to be processed by this
///
/// SAFETY: `oc` must be an index that's at most `out_channels - simd_lanes`, so the full vector
/// is in bounds. Weights and `out` must be channels last (with `stride == 1`).
unsafe fn conv2d_remainder<S: Simd, E: VMulAdd>(
    x: ArrayView3<E>,
    weights: ArrayView4<E>,
    out: &mut ArrayViewMut3<E>,
    bias: Vector<S, E>,
    oc: usize,
    ic_off: usize,
    owb_end: usize,
    stride_h: usize,
    stride_w: usize,
    dilate_h: usize,
    dilate_w: usize,
    pad_h: usize,
    pad_w: usize,
    k_height: usize,
    k_width: usize,
) {
// Load a full vector from the weights. This is guaranteed to be in bounds
// as long as `oc <= out_channels - simd_lanes` and out channels are last.
// We need to ensure the weights are reshaped appropriately.
// The loop bounds ensure `ic`, `ih` and `iw` are always in bounds, but the
// compiler can't prove this. We can't use `as_slice` with fixed bounds
// because we want to support arbitrary input layouts. So an unchecked load
// is used.
// Store a full vector from the output. This is guaranteed to be in bounds
// as long as `oc <= out_channels - simd_lanes` and oc stride is 1. We create `out` with
// channels last, so this always holds.
// Load a full vector from the weights. This is guaranteed to be in bounds
// as long as `oc <= out_channels - simd_lanes` and out channels are last.
// We need to ensure the weights are reshaped appropriately.
// The loop bounds ensure `ic`, `ih` and `iw` are always in bounds, but the
// compiler can't prove this. We can't use `as_slice` with fixed bounds
// because we want to support arbitrary input layouts. So an unchecked load
// is used.
// Store a full vector from the output. This is guaranteed to be in bounds
// as long as `oc <= out_channels - simd_lanes` and oc stride is 1. We create `out` with
// channels last, so this always holds.
/// Execute the unrolled and unpadded portion of the convolution. Any pixel that is more than
/// `pad_h` away from the horizontal border, and `pad_w` away from the vertical border is
/// guaranteed to always be in bounds (because of the way out size is calculated).
///
/// SAFETY: `oc` must be an index that's at most `out_channels - simd_lanes`, so the full vector
/// is in bounds. Weights and `out` must be channels last (with `stride == 1`).
        unsafe fn conv2d_inner_nopad<S: Simd, E: VMulAdd>(
            x: &ArrayView3<E>,
            weights: &ArrayView4<E>,
            out: &mut ArrayViewMut2<E>,
            bias: Vector<S, E>,
            oh: usize,
            ow: usize,
            oc: usize,
            ic_off: usize,
            stride_h: usize,
            stride_w: usize,
            dilate_h: usize,
            dilate_w: usize,
            k_height: usize,
            k_width: usize,
            pad_h: usize,
            pad_w: usize,
        ) {
// Load a full vector from the weights. This is guaranteed to be in bounds
// as long as `oc <= out_channels - simd_lanes` and out channels are last.
// We need to ensure the weights are reshaped appropriately.
// The loop bounds ensure `ic`, `ih` and `iw` are always in bounds, but the
// compiler can't prove this. We can't use `as_slice` with fixed bounds
// because we want to support arbitrary input layouts. So an unchecked load
// is used.
// Store a full vector from the output. This is guaranteed to be in bounds
// as long as `oc <= out_channels - simd_lanes` and oc stride is 1. We create `out` with
// channels last, so this always holds.
/// Execute the unrolled and unpadded portion of the convolution. Any pixel that is more than
/// `pad_h` away from the horizontal border, and `pad_w` away from the vertical border is
/// guaranteed to always be in bounds (because of the way out size is calculated).
///
/// SAFETY: `oc` must be an index that's at most `out_channels - simd_lanes`, so the full vector
/// is in bounds. Weights and `out` must be channels last (with `stride == 1`).
        unsafe fn conv2d_inner_nopad_nostride<S: Simd, E: VMulAdd>(
            x: &ArrayView3<E>,
            weights: &ArrayView4<E>,
            out: &mut ArrayViewMut2<E>,
            bias: Vector<S, E>,
            oh: usize,
            ow: usize,
            oc: usize,
            ic_off: usize,
            k_height: usize,
            k_width: usize,
            pad_h: usize,
            pad_w: usize,
        ) {
// Load a full vector from the weights. This is guaranteed to be in bounds
// as long as `oc <= out_channels - simd_lanes` and out channels are last.
// We need to ensure the weights are reshaped appropriately.
// The loop bounds ensure `ic`, `ih` and `iw` are always in bounds, but the
// compiler can't prove this. We can't use `as_slice` with fixed bounds
// because we want to support arbitrary input layouts. So an unchecked load
// is used.
// Store a full vector from the output. This is guaranteed to be in bounds
// as long as `oc <= out_channels - simd_lanes` and oc stride is 1. We create `out` with
// channels last, so this always holds.
fn is_accelerated_impl<S: Simd, T: VOrd>(_x: PhantomData<T>) -> bool {
fn is_accelerated<T: VOrd>() -> bool {
pub(crate) fn try_max_pool2d_simd<E: Element>(
    x: NdArrayTensor<E>,
    ksize: [usize;
fn cast<T, E>(tensor: NdArrayTensor<T>) -> NdArrayTensor<E> {
// Until you can use associated constants as array size, we need to hardcode this.
// The most common config (x86-v3) has 16 registers, so use half of them for accumulators.
    const BLOCK_REGISTERS: usize = 8;
    pub(crate) fn max_pool2d_nhwc<E: Element + VOrd + MinMax>(
        x: NdArrayTensor<E>,
        kernel_size: [usize;
// Floor division ensures `blocks * lanes * blocking factor` is always `<= out_channels`.
// An exclusive loop will always have `lanes * blocking factor` elements in bounds.
// Floor division means simd_end is always divisible by `lanes` and `<= out_channels`. An
// exclusive loop will always have `lanes` elements in bounds.
// SAFETY: Loop ranges are non-overlapping, so the unsafe shared reference is safe.
// SAFETY: See `loop_unblocked`
// SAFETY: Loop ranges are non-overlapping, so the unsafe shared reference is safe.
/// Execute the blocked (unrolled) portion of the pool.
    fn loop_blocked<'a, S: Simd, E: Element + VOrd + MinMax>(
        x: ArrayView3<'a, E>,
        mut out: ArrayViewMut3<'a, E>,
        kernel_size: [usize;
// If outside padding area, kernels are guaranteed to be in bounds
// SAFETY:
// Load a full vector from x[N * lanes]. This is bounds checked by the
// slice above.
// SAFETY:
// Store a full vector to out[N * lanes]. This is bounds checked by the
// slice above.
// Border pixels need bounds checks
// SAFETY:
// Load a full vector from x[N * lanes]. This is bounds checked by the
// slice above.
// SAFETY:
// Store a full vector to out[N * lanes]. This is bounds checked by the
// slice above.
/// Execute the unblocked (not unrolled) portion of the pool.
///
/// SAFETY: Safe as long as `ch + simd_lanes <= out_channels`.
    unsafe fn loop_unblocked<'a, S: Simd, E: Element + VOrd + MinMax>(
        x: ArrayView3<'a, E>,
        mut out: ArrayViewMut3<'a, E>,
        kernel_size: [usize;
// Load a full vector from `x`. In bounds as long as `out_channels >= ch + lanes`
// Store a full vector to `out`. In bounds as long as `out_channels >= ch + lanes`.
// Border pixels need bounds checks
// Load a full vector from `x`. In bounds as long as `out_channels >= ch + lanes`
// Store a full vector to `out`. In bounds as long as `out_channels >= ch + lanes`.
    fn loop_scalar<E: Element + MinMax>(
        x: ArrayView3<'_, E>,
        mut out: ArrayViewMut3<'_, E>,
        kernel_size: [usize;
pub trait SimdUnop<T: Scalar, Out: Scalar> {
    fn apply_vec<S: Simd>(input: Vector<S, T>) -> Vector<S, Out>;
    fn apply(input: T) -> Out;
    fn is_accelerated<S: Simd>() -> bool;
pub struct RecipVec;
impl SimdUnop<f32, f32> for RecipVec {
    fn apply_vec<S: Simd>(input: Vector<S, f32>) -> Vector<S, f32> {
    fn apply(input: f32) -> f32 {
    fn is_accelerated<S: Simd>() -> bool {
pub struct VecAbs;
impl<T: VAbs + Signed> SimdUnop<T, T> for VecAbs {
    fn apply_vec<S: Simd>(input: Vector<S, T>) -> Vector<S, T> {
    fn apply(input: T) -> T {
    fn is_accelerated<S: Simd>() -> bool {
pub struct VecBitNot;
impl<T: VBitNot> SimdUnop<T, T> for VecBitNot {
    fn apply_vec<S: Simd>(input: Vector<S, T>) -> Vector<S, T> {
    fn apply(input: T) -> T {
    fn is_accelerated<S: Simd>() -> bool {
fn is_accelerated<S: Simd, T: Scalar, Out: Scalar, Op: SimdUnop<T, Out>>(
    _x: PhantomData<(T, Out, Op)>,
) -> bool {
pub fn try_unary_simd<
    E: NdArrayElement,
    EOut: NdArrayElement,
    T: NdArrayElement + Scalar,
    Out: NdArrayElement + Scalar,
    Op: SimdUnop<T, Out>,
>(
    input: NdArrayTensor<E>,
) -> Result<NdArrayTensor<EOut>, NdArrayTensor<E>> {
// Used to assert traits based on the dynamic `DType`.
// Used to assert traits based on the dynamic `DType`.
/// Execute operation in line.
/// SAFETY:
/// Must ensure `size_of::<T> == size_of::<Out>` and `align_of::<T> >= align_of::<Out>`.
unsafe fn unary_scalar_simd_inplace<
    T: NdArrayElement + Scalar,
    Out: NdArrayElement + Scalar,
    Op: SimdUnop<T, Out>,
>(
    input: NdArrayTensor<T>,
) -> NdArrayTensor<Out> {
// This is only called when in and out have the same size, so it's safe
// Buffer has the same elem size and is filled with the operation output, so this is safe
fn unary_scalar_simd_owned<
    T: NdArrayElement + Scalar,
    Out: NdArrayElement + Scalar,
    Op: SimdUnop<T, Out>,
>(
    input: NdArrayTensor<T>,
) -> NdArrayTensor<Out> {
fn unary_slice<
    'a,
    S: Simd,
    T: NdArrayElement + Scalar,
    Out: NdArrayElement + Scalar,
    Op: SimdUnop<T, Out>,
>(
    input: &'a [T],
    out: &'a mut [Out],
    _op: PhantomData<Op>,
) where
    'a: 'a,
{
// Load one full vector from `input`.
// SAFETY: Guaranteed to be in bounds because `len == 8 * lanes`
// Store one full vector to `out`.
// SAFETY: Guaranteed to be in bounds because `len == 8 * lanes`
// Load one full vector from `input`.
// SAFETY: Guaranteed to be in bounds because `len == lanes`
// Store one full vector to `out`.
// SAFETY: Guaranteed to be in bounds because `len == lanes`
/// Execute operation in line.
/// SAFETY:
/// Must ensure `size_of::<T> == size_of::<Out>` and `align_of::<T> >= align_of::<Out>`.
unsafe fn unary_slice_inplace<
    'a,
    S: Simd,
    T: NdArrayElement + Scalar,
    Out: NdArrayElement + Scalar,
    Op: SimdUnop<T, Out>,
>(
    buf: &'a mut [T],
    _op: PhantomData<(Out, Op)>,
) where
    'a: 'a,
{
// Load a full vector from the aligned portion of the buffer.
// SAFETY: `align_to_mut` guarantees we're aligned to `T::Vector`'s size, and there is
// always a full vector in bounds.
            let s~N = unsafe {
// Store a full vector at the same position as the input. Cast is safe because `Out` is
// size and align compatible
// Load a full vector from the aligned portion of the buffer.
// SAFETY: `align_to_mut` guarantees we're aligned to `T::Vector`'s size, and there is
// always a full vector in bounds.
        let s0 = unsafe {
// Store a full vector at the same position as the input. Cast is safe because `Out` is
// size and align compatible
// Language
// Current crate
// Workspace crates
fn round_ties_even_wrapper(x: f64) -> f64 {
fn round_ties_even_wrapper(x: f64) -> f64 {
impl<E: FloatNdArrayElement, I: IntNdArrayElement, Q: QuantElement> FloatTensorOps<Self>
    for NdArray<E, I, Q>
{
    fn float_from_data(data: TensorData, _device: &NdArrayDevice) -> FloatTensor<Self> {
    fn float_random(
        shape: Shape,
        distribution: Distribution,
        device: &NdArrayDevice,
    ) -> FloatTensor<Self> {
    async fn float_into_data(tensor: FloatTensor<Self>) -> TensorData {
    fn float_device(_tensor: &FloatTensor<Self>) -> NdArrayDevice {
    fn float_to_device(tensor: FloatTensor<Self>, _device: &NdArrayDevice) -> FloatTensor<Self> {
    fn float_empty(shape: Shape, device: &<NdArray<E> as Backend>::Device) -> FloatTensor<Self> {
    fn float_add(lhs: FloatTensor<Self>, rhs: FloatTensor<Self>) -> FloatTensor<Self> {
    fn float_add_scalar(lhs: FloatTensor<Self>, rhs: E) -> FloatTensor<Self> {
    fn float_sub(lhs: FloatTensor<Self>, rhs: FloatTensor<Self>) -> FloatTensor<Self> {
    fn float_sub_scalar(lhs: FloatTensor<Self>, rhs: E) -> FloatTensor<Self> {
    fn float_mul(lhs: FloatTensor<Self>, rhs: FloatTensor<Self>) -> FloatTensor<Self> {
    fn float_mul_scalar(lhs: FloatTensor<Self>, rhs: E) -> FloatTensor<Self> {
    fn float_div(lhs: FloatTensor<Self>, rhs: FloatTensor<Self>) -> FloatTensor<Self> {
    fn float_div_scalar(lhs: FloatTensor<Self>, rhs: E) -> FloatTensor<Self> {
    fn float_remainder(lhs: FloatTensor<Self>, rhs: FloatTensor<Self>) -> FloatTensor<Self> {
    fn float_remainder_scalar(lhs: FloatTensor<Self>, rhs: E) -> FloatTensor<Self> {
    fn float_matmul(lhs: FloatTensor<Self>, rhs: FloatTensor<Self>) -> FloatTensor<Self> {
    fn float_neg(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
    fn float_recip(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
    fn float_swap_dims(tensor: FloatTensor<Self>, dim1: usize, dim2: usize) -> FloatTensor<Self> {
    fn float_reshape(tensor: FloatTensor<Self>, shape: Shape) -> FloatTensor<Self> {
    fn float_gather(
        dim: usize,
        tensor: FloatTensor<Self>,
        indices: NdArrayTensor<I>,
    ) -> FloatTensor<Self> {
    fn float_scatter(
        dim: usize,
        tensor: FloatTensor<Self>,
        indices: NdArrayTensor<I>,
        value: FloatTensor<Self>,
    ) -> FloatTensor<Self> {
    fn float_select(
        tensor: FloatTensor<Self>,
        dim: usize,
        indices: NdArrayTensor<I>,
    ) -> FloatTensor<Self> {
    fn float_select_assign(
        tensor: FloatTensor<Self>,
        dim: usize,
        indices: NdArrayTensor<I>,
        value: FloatTensor<Self>,
    ) -> FloatTensor<Self> {
    fn float_slice(tensor: FloatTensor<Self>, ranges: &[Range<usize>]) -> FloatTensor<Self> {
    fn float_slice_assign(
        tensor: FloatTensor<Self>,
        ranges: &[Range<usize>],
        value: FloatTensor<Self>,
    ) -> FloatTensor<Self> {
    fn float_mask_where(
        tensor: FloatTensor<Self>,
        mask: NdArrayTensor<bool>,
        value: FloatTensor<Self>,
    ) -> FloatTensor<Self> {
    fn float_mask_fill(
        tensor: FloatTensor<Self>,
        mask: NdArrayTensor<bool>,
        value: E,
    ) -> FloatTensor<Self> {
    fn float_equal(lhs: FloatTensor<Self>, rhs: FloatTensor<Self>) -> NdArrayTensor<bool> {
    fn float_equal_elem(lhs: FloatTensor<Self>, rhs: E) -> NdArrayTensor<bool> {
    fn float_greater(lhs: FloatTensor<Self>, rhs: FloatTensor<Self>) -> NdArrayTensor<bool> {
    fn float_greater_elem(lhs: FloatTensor<Self>, rhs: E) -> NdArrayTensor<bool> {
    fn float_greater_equal(lhs: FloatTensor<Self>, rhs: FloatTensor<Self>) -> NdArrayTensor<bool> {
    fn float_greater_equal_elem(lhs: FloatTensor<Self>, rhs: E) -> NdArrayTensor<bool> {
    fn float_lower(lhs: FloatTensor<Self>, rhs: FloatTensor<Self>) -> NdArrayTensor<bool> {
    fn float_lower_elem(lhs: FloatTensor<Self>, rhs: E) -> NdArrayTensor<bool> {
    fn float_lower_equal(lhs: FloatTensor<Self>, rhs: FloatTensor<Self>) -> NdArrayTensor<bool> {
    fn float_lower_equal_elem(lhs: FloatTensor<Self>, rhs: E) -> NdArrayTensor<bool> {
    fn float_detach(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
    fn float_mean(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
    fn float_sum(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
    fn float_mean_dim(tensor: FloatTensor<Self>, dim: usize) -> FloatTensor<Self> {
    fn float_sum_dim(tensor: FloatTensor<Self>, dim: usize) -> FloatTensor<Self> {
    fn float_argmax(tensor: FloatTensor<Self>, dim: usize) -> NdArrayTensor<I> {
    fn float_argmin(tensor: FloatTensor<Self>, dim: usize) -> NdArrayTensor<I> {
    fn float_exp(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
    fn float_log(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
    fn float_prod(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
    fn float_prod_dim(tensor: FloatTensor<Self>, dim: usize) -> FloatTensor<Self> {
    fn float_log1p(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
    fn float_powf_scalar(tensor: FloatTensor<Self>, value: f32) -> FloatTensor<Self> {
// Happens often and is faster.
// Is faster then powf
// Default
    fn float_sqrt(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
    fn float_abs(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
    fn float_cos(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
    fn float_sin(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
    fn float_tanh(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
    fn float_round(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
    fn float_floor(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
    fn float_ceil(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
    fn float_erf(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
    fn float_cat(tensors: Vec<FloatTensor<Self>>, dim: usize) -> FloatTensor<Self> {
    fn float_clamp_min(tensor: FloatTensor<Self>, min: E) -> FloatTensor<Self> {
    fn float_clamp_max(tensor: FloatTensor<Self>, max: E) -> FloatTensor<Self> {
    fn float_clamp(tensor: FloatTensor<Self>, min: E, max: E) -> FloatTensor<Self> {
    fn float_into_int(tensor: FloatTensor<Self>) -> NdArrayTensor<I> {
    fn float_powf(lhs: FloatTensor<Self>, rhs: FloatTensor<Self>) -> FloatTensor<Self> {
    fn float_permute(tensor: FloatTensor<Self>, axes: &[usize]) -> FloatTensor<Self> {
    fn float_flip(tensor: FloatTensor<Self>, axes: &[usize]) -> FloatTensor<Self> {
    fn float_sign(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
    fn float_expand(tensor: FloatTensor<Self>, shape: Shape) -> FloatTensor<Self> {
    fn float_cast(tensor: FloatTensor<Self>, dtype: FloatDType) -> FloatTensor<Self> {
        fn cast<E1: FloatNdArrayElement, E2: FloatNdArrayElement>(
            tensor: &NdArrayTensor<E1>,
        ) -> NdArrayTensor<E2> {
// No cast
// F32 to F64
// F64 to F32
impl<E: FloatNdArrayElement, I: IntNdArrayElement, Q: QuantElement> TransactionOps<Self>
    for NdArray<E, I, Q>
{
/// Similar to `SyncUnsafeCell` see [Rust issues](https://github.com/rust-lang/rust/issues/95439).
pub(crate) struct UnsafeSharedRef<'a, T> {
unsafe impl<T> Sync for UnsafeSharedRef<'_, T> {
impl<'a, T> UnsafeSharedRef<'a, T> {
    pub fn new(data: &'a mut T) -> Self {
    pub unsafe fn get(&self) -> &'a mut T {
/// Tensor primitive used by the [ndarray backend](crate::NdArray).
pub struct NdArrayTensor<E> {
/// Dynamic array that contains the data of type E.
impl<E: Element> TensorMetadata for NdArrayTensor<E> {
    fn dtype(&self) -> DType {
    fn shape(&self) -> Shape {
/// Float tensor primitive.
pub enum NdArrayTensorFloat {
/// 32-bit float.
/// 64-bit float.
impl From<NdArrayTensor<f32>> for NdArrayTensorFloat {
    fn from(value: NdArrayTensor<f32>) -> Self {
impl From<NdArrayTensor<f64>> for NdArrayTensorFloat {
    fn from(value: NdArrayTensor<f64>) -> Self {
impl TensorMetadata for NdArrayTensorFloat {
    fn dtype(&self) -> DType {
    fn shape(&self) -> Shape {
/// Macro to create a new [float tensor](NdArrayTensorFloat) based on the element type.
// Op executed with default dtype
// FloatNdArrayElement only implemented for f64 and f32
/// Macro to execute an operation a given element type.
///
/// # Panics
/// Since there is no automatic type cast at this time, binary operations for different
/// floating point precision data types will panic with a data type mismatch.
// Binary op: type automatically inferred by the compiler
// Binary op: generic type cannot be inferred for an operation
                type $element = f64;
                type $element = f32;
// Binary op: type automatically inferred by the compiler but return type is not a float tensor
// Unary op: type automatically inferred by the compiler
// Unary op: generic type cannot be inferred for an operation
                type $element = f64;
                type $element = f32;
// Unary op: type automatically inferred by the compiler but return type is not a float tensor
// Unary op: generic type cannot be inferred for an operation and return type is not a float tensor
                type $element = f64;
                type $element = f32;
    impl<E> NdArrayTensor<E>
    where
        E: Element,
    {
        pub(crate) fn into_data(self) -> TensorData {
        pub(crate) fn is_contiguous(&self) -> bool {
/// Converts a slice of usize to a typed dimension.
/// Reshapes an array into a tensor.
impl<E> NdArrayTensor<E>
where
    E: Element,
{
/// Create a new [ndarray tensor](NdArrayTensor) from [data](TensorData).
    pub fn from_data(mut data: TensorData) -> NdArrayTensor<E> {
// Safety: TensorData checks shape validity on creation, so we don't need to repeat that check here
/// A quantized tensor for the ndarray backend.
pub struct NdArrayQTensor<Q: QuantElement> {
/// The quantized tensor.
/// The quantization scheme.
/// The quantization parameters.
impl<Q: QuantElement> NdArrayQTensor<Q> {
/// Returns the quantization strategy, including quantization parameters, for the given tensor.
    pub fn strategy(&self) -> QuantizationStrategy {
impl<Q: QuantElement> QTensorPrimitive for NdArrayQTensor<Q> {
    fn scheme(&self) -> &QuantScheme {
impl<Q: QuantElement> TensorMetadata for NdArrayQTensor<Q> {
    fn dtype(&self) -> DType {
    fn shape(&self) -> Shape {
    fn should_support_into_and_from_data_1d() {
    fn should_support_into_and_from_data_2d() {
    fn should_support_into_and_from_data_3d() {
    fn should_support_into_and_from_data_4d() {
    fn should_support_qtensor_strategy() {
        type B = NdArray<f32, i64, i8>;
/// [`Activation`] Configuration.
pub enum ActivationConfig {
/// [`Gelu`] activation layer.
/// [`PRelu`] activation layer.
/// [`Relu`] activation layer.
/// [`LeakyRelu`] activation layer.
/// [`SwiGlu`] activation layer.
/// [`Sigmoid`] activation layer.
/// [`Tanh`] activation layer.
/// [`HardSigmoid`] activation layer.
impl From<PReluConfig> for ActivationConfig {
    fn from(config: PReluConfig) -> Self {
impl From<LeakyReluConfig> for ActivationConfig {
    fn from(config: LeakyReluConfig) -> Self {
impl From<SwiGluConfig> for ActivationConfig {
    fn from(config: SwiGluConfig) -> Self {
impl From<HardSigmoidConfig> for ActivationConfig {
    fn from(config: HardSigmoidConfig) -> Self {
impl ActivationConfig {
/// Initialize a wrapped activation layer.
    pub fn init<B: Backend>(&self, device: &B::Device) -> Activation<B> {
/// Activation Layer Wrapper.
///
/// Provides support for many in-built `burn::nn` activations.
pub enum Activation<B: Backend> {
/// [`Gelu`] activation layer.
/// [`PRelu`] activation layer.
/// [`Relu`] activation layer.
/// [`LeakyRelu`] activation layer.
/// [`SwiGlu`] activation layer.
/// [`Sigmoid`] activation layer.
/// [`Tanh`] activation layer.
/// [`HardSigmoid`] activation layer.
impl<B: Backend> From<Gelu> for Activation<B> {
    fn from(layer: Gelu) -> Self {
impl<B: Backend> From<PRelu<B>> for Activation<B> {
    fn from(layer: PRelu<B>) -> Self {
impl<B: Backend> From<Relu> for Activation<B> {
    fn from(layer: Relu) -> Self {
impl<B: Backend> From<LeakyRelu> for Activation<B> {
    fn from(layer: LeakyRelu) -> Self {
impl<B: Backend> From<SwiGlu<B>> for Activation<B> {
    fn from(layer: SwiGlu<B>) -> Self {
impl<B: Backend> From<Sigmoid> for Activation<B> {
    fn from(layer: Sigmoid) -> Self {
impl<B: Backend> From<Tanh> for Activation<B> {
    fn from(layer: Tanh) -> Self {
impl<B: Backend> From<HardSigmoid> for Activation<B> {
    fn from(layer: HardSigmoid) -> Self {
impl<B: Backend> Activation<B> {
/// Forward pass.
    pub fn forward<const D: usize>(&self, input: Tensor<B, D>) -> Tensor<B, D> {
    fn make_input<B: Backend>(device: &B::Device) -> Tensor<B, 2> {
    fn expect_tensor<B: Backend, const D: usize>(actual: Tensor<B, D>, expected: Tensor<B, D>) {
    fn check_stateless_config_output<B: Backend, const D: usize>(
        config: ActivationConfig,
        input: Tensor<B, D>,
        expected: Tensor<B, D>,
        device: &B::Device,
    ) {
    fn test_gelu() {
    fn test_prelu() {
    fn test_relu() {
    fn test_leaky_relu() {
    fn test_swi_glu() {
// Clone the initialized weights.
    fn test_sigmoid() {
    fn test_tanh() {
    fn test_hard_sigmoid() {
/// Applies the Gaussian Error Linear Units function element-wise.
/// See also [gelu](burn::tensor::activation::gelu)
pub struct Gelu;
impl Gelu {
/// Create the module.
    pub fn new() -> Self {
/// Applies the forward pass on the input tensor.
///
/// # Shapes
///
/// - input: `[..., any]`
/// - output: `[..., any]`
    pub fn forward<B: Backend, const D: usize>(&self, input: Tensor<B, D>) -> Tensor<B, D> {
    fn display() {
/// Applies the gated linear unit function.
///
/// See also [glu](burn::tensor::activation::glu)
pub struct GLU {
impl GLU {
/// Create the module.
///
/// # Arguments
/// * `dim` - The dimension on which to split the input.
    pub fn new(dim: usize) -> Self {
/// Applies the gated linear unit function.
///
/// GLU(a,b)=a⊗σ(b) where `a` is the first half of the input matrices and `b` is the second half.
///
/// **Note**:
/// * The size of the input tensor along `dim` must be divisible by 2.
///
/// ### Arguments
/// * `tensor` - The input tensor.
///
/// ### Returns
/// * A tensor with the same shape as the input, except the size along `dim` is halved.
    pub fn forward<B: Backend, const D: usize>(&self, input: Tensor<B, D>) -> Tensor<B, D> {
    fn display() {
/// Hard Sigmoid layer.
///
/// Should be created with [HardSigmoidConfig](HardSigmoidConfig).
pub struct HardSigmoid {
/// The alpha value.
/// The beta value.
/// Configuration to create a [Hard Sigmoid](HardSigmoid) layer using the [init function](HardSigmoidConfig::init).
pub struct HardSigmoidConfig {
/// The alpha value. Default is 0.2
/// The beta value. Default is 0.5
impl HardSigmoidConfig {
/// Initialize a new [Hard Sigmoid](HardSigmoid) Layer
    pub fn init(&self) -> HardSigmoid {
impl ModuleDisplay for HardSigmoid {
    fn custom_settings(&self) -> Option<DisplaySettings> {
    fn custom_content(&self, content: Content) -> Option<Content> {
impl HardSigmoid {
/// Forward pass for the Hard Sigmoid layer.
///
/// See [hard_sigmoid](burn::tensor::activation::hard_sigmoid) for more information.
///
/// # Shapes
/// - input: `[..., any]`
/// - output: `[..., any]`
    pub fn forward<B: Backend, const D: usize>(&self, input: Tensor<B, D>) -> Tensor<B, D> {
    type FT = FloatElem<TestBackend>;
    fn test_hard_sigmoid_forward() {
    fn display() {
/// Leaky ReLu layer.
///
/// Should be created with [LeakyReluConfig](LeakyReluConfig).
pub struct LeakyRelu {
/// The negative slope.
/// Configuration to create a [Leaky Relu](LeakyRelu) layer using the [init function](LeakyReluConfig::init).
pub struct LeakyReluConfig {
/// The negative slope. Default is 0.01
impl LeakyReluConfig {
/// Initialize a new [Leaky Relu](LeakyRelu) Layer
    pub fn init(&self) -> LeakyRelu {
impl ModuleDisplay for LeakyRelu {
    fn custom_settings(&self) -> Option<DisplaySettings> {
    fn custom_content(&self, content: Content) -> Option<Content> {
impl LeakyRelu {
/// Forward pass for the Leaky ReLu layer.
///
/// See [leaky_relu](burn::tensor::activation::leaky_relu) for more information.
///
/// # Shapes
/// - input: `[..., any]`
/// - output: `[..., any]`
    pub fn forward<B: Backend, const D: usize>(&self, input: Tensor<B, D>) -> Tensor<B, D> {
    type FT = FloatElem<TestBackend>;
    fn test_leaky_relu_forward() {
    fn test_leaky_relu_forward_multi_dim() {
    fn display() {
//! # Activation Layers
//!
//! Users who desire a selectable activation function should
//! consider [`Activation`], which provides an abstraction over:
//! * [`Relu`] - the default,
//! * ['PRelu']
//! * [`Gelu`]
//! * [`LeakyRelu`]
//! * [`SwiGlu`]
//! * [`Sigmoid`]
//! * [`HardSigmoid`]
//! * [`Tanh`]
//!
//! The activation layer [`GLU`] has shape-changing behaviors
//! not compatible with the common API, and is not included
//! in the abstraction wrappers.
// These are pub(crate) for dual-export in `nn` without re-exporting
// all of `nn.activation`, or manually listing each symbol.
/// Parametric Relu layer.
///
/// Should be created using [PReluConfig]
pub struct PRelu<B: Backend> {
/// the weights learnt for PReLu. can be of shape \[1\] or \[num_parameters\] in which case it must
/// be the same as number of channels in the input tensor
/// Alpha value for the PRelu layer
impl<B: Backend> ModuleDisplay for PRelu<B> {
    fn custom_settings(&self) -> Option<DisplaySettings> {
    fn custom_content(&self, content: Content) -> Option<Content> {
/// Configuration to create a [Parametric Relu](PRelu) layer using the [init function](PReluConfig::init).
pub struct PReluConfig {
/// The number of parameters.
/// The learnable weight alpha. Default is 0.25
impl PReluConfig {
/// Initialize a new [Parametric Relu](PRelu) Layer
    pub fn init<B: Backend>(&self, device: &B::Device) -> PRelu<B> {
// alpha is a tensor of length num_parameters
impl<B: Backend> PRelu<B> {
/// Applies the forward pass on the input tensor.
///
/// # Shapes
///
/// - input: `[..., any]`
/// - output: `[..., any]`
///
/// See also [prelu](burn::tensor::activation::prelu) for more information.
    pub fn forward<const D: usize>(&self, input: Tensor<B, D>) -> Tensor<B, D> {
    fn display() {
/// Applies the rectified linear unit function element-wise
/// See also [relu](burn::tensor::activation::relu)
///
pub struct Relu;
impl Relu {
/// Create the module.
    pub fn new() -> Self {
/// Applies the forward pass on the input tensor.
///
/// # Shapes
///
/// - input: `[..., any]`
/// - output: `[..., any]`
    pub fn forward<B: Backend, const D: usize>(&self, input: Tensor<B, D>) -> Tensor<B, D> {
    fn display() {
/// Applies the sigmoid function element-wise
/// See also [sigmoid](burn::tensor::activation::sigmoid)
pub struct Sigmoid;
impl Sigmoid {
/// Create the module.
    pub fn new() -> Self {
/// Applies the forward pass on the input tensor.
///
/// # Shapes
///
/// - input: `[..., any]`
/// - output: `[..., any]`
    pub fn forward<B: Backend, const D: usize>(&self, input: Tensor<B, D>) -> Tensor<B, D> {
    fn display() {
/// Configuration to create a [SwiGlu](SwiGlu) activation layer using the [init function](SwiGluConfig::init).
pub struct SwiGluConfig {
/// The size of the input features.
/// The size of the output features.
/// If a bias should be applied during the linear transformation. Default behaviour is False
/// for SwiGLU activation implementations.
/// The type of function used to initialize the linear layer parameters
/// The layout in which the linear parameters are stored.
/// Applies the SwiGLU or Swish Gated Linear Unit to the input tensor.
/// The SwiGLU activation function is defined as:
/// `SwiGLU(x) = Swish(W_inner * x + b_inner) * (W_outer * x + b_outer)`
///
/// Should be created with [SwiGluConfig].
pub struct SwiGlu<B: Backend> {
/// The inner linear layer for Swish activation function
/// with `d_input` input features and `d_output` output features.
/// The outer linear layer for element wise multiplication
/// with `d_input` input features and `d_output` output features.
impl<B: Backend> ModuleDisplay for SwiGlu<B> {
    fn custom_settings(&self) -> Option<DisplaySettings> {
    fn custom_content(&self, content: Content) -> Option<Content> {
impl SwiGluConfig {
/// Initialize a new [SwiGLU](SwiGlu) activation layer.
    pub fn init<B: Backend>(&self, device: &B::Device) -> SwiGlu<B> {
impl<B: Backend> SwiGlu<B> {
/// Applies the Swish Gated Linear Unit to the input tensor.
///
/// # Shapes
///
/// - input: `[batch_size, seq_length, d_input]`
/// - output: `[batch_size, seq_length, d_output]`
    pub fn forward<const D: usize>(&self, input: Tensor<B, D>) -> Tensor<B, D> {
    type FT = FloatElem<TestBackend>;
    fn test_swiglu_forward_no_bias() {
    fn test_swiglu_forward_with_bias() {
    fn display() {
/// Applies the tanh activation function element-wise
/// See also [tanh](burn::tensor::activation::tanh)
pub struct Tanh;
impl Tanh {
/// Create the module.
    pub fn new() -> Self {
/// Applies the forward pass on the input tensor.
///
/// # Shapes
///
/// - input: `[..., any]`
/// - output: `[..., any]`
    pub fn forward<B: Backend, const D: usize>(&self, input: Tensor<B, D>) -> Tensor<B, D> {
    fn display() {
//! Burn neural network module.
/// Loss module
/// Neural network modules implementations.
// For backward compat, `burn::nn::Initializer`
/// Backend for test cases
pub type TestBackend = burn_ndarray::NdArray<f32>;
/// Backend for test cases
pub type TestBackend = burn_tch::LibTorch<f32>;
/// Backend for test cases
pub type TestBackend = burn_wgpu::Wgpu;
/// Backend for test cases
pub type TestBackend = burn_cuda::Cuda;
/// Backend for test cases
pub type TestBackend = burn_rocm::Rocm;
/// Backend for autodiff test cases
pub type TestAutodiffBackend = burn_autodiff::Autodiff<TestBackend>;
/// Configuration to create a [Binary Cross-entropy loss](BinaryCrossEntropyLoss) using the [init function](BinaryCrossEntropyLossConfig::init).
pub struct BinaryCrossEntropyLossConfig {
/// Create weighted binary cross-entropy with a weight for each class.
///
/// The loss of a specific sample will simply be multiplied by its label weight.
/// Create binary cross-entropy with label smoothing according to [When Does Label Smoothing Help?](https://arxiv.org/abs/1906.02629).
///
/// Hard labels {0, 1} will be changed to `y_smoothed = y(1 - a) + a / num_classes`.
/// Alpha = 0 would be the same as default.
/// Treat the inputs as logits, applying a sigmoid activation when computing the loss.
impl BinaryCrossEntropyLossConfig {
/// Initialize [Binary Cross-entropy loss](BinaryCrossEntropyLoss).
    pub fn init<B: Backend>(&self, device: &B::Device) -> BinaryCrossEntropyLoss<B> {
    fn assertions(&self) {
/// Calculate the binary cross entropy loss from the input logits and the targets.
///
/// Should be created using [BinaryCrossEntropyLossConfig]
pub struct BinaryCrossEntropyLoss<B: Backend> {
/// Weights for cross-entropy.
/// Label smoothing alpha.
/// Treat the inputs as logits
impl<B: Backend> ModuleDisplay for BinaryCrossEntropyLoss<B> {
    fn custom_settings(&self) -> Option<DisplaySettings> {
    fn custom_content(&self, content: Content) -> Option<Content> {
impl<B: Backend> BinaryCrossEntropyLoss<B> {
/// Compute the criterion on the input tensor.
///
/// # Shapes
///
/// Binary:
/// - logits: `[batch_size]`
/// - targets: `[batch_size]`
///
/// Multi-label:
/// - logits: `[batch_size, num_classes]`
/// - targets: `[batch_size, num_classes]`
    pub fn forward<const D: usize>(
        &self,
        logits: Tensor<B, D>,
        targets: Tensor<B, D, Int>,
    ) -> Tensor<B, 1> {
// Numerically stable by combining `log(sigmoid(x))` with `log_sigmoid(x)`
// - (target * log(input) + (1 - target) * log(1 - input))
// https://github.com/tracel-ai/burn/issues/2739: clamp at -100.0 to avoid undefined values
// Flatten targets and expand resulting weights to make it compatible with
// Tensor<B, D> for binary 1-D case
    fn assertions<const D: usize>(&self, logits: &Tensor<B, D>, targets: &Tensor<B, D, Int>) {
    type FT = FloatElem<TestBackend>;
    fn test_binary_cross_entropy_preds_all_correct() {
    fn test_binary_cross_entropy_preds_all_incorrect() {
// clamped value
    fn test_binary_cross_entropy() {
// import torch
// from torch import nn
// input = torch.tensor([0.8271, 0.9626, 0.3796, 0.2355])
// target = torch.tensor([0., 1., 0., 1.])
// loss = nn.BCELoss()
// sigmoid = nn.Sigmoid()
// out = loss(sigmoid(input), target) # tensor(0.7491)
    fn test_binary_cross_entropy_with_logits() {
    fn test_binary_cross_entropy_with_weights() {
// import torch
// from torch import nn
// input = torch.tensor([0.8271, 0.9626, 0.3796, 0.2355])
// target = torch.tensor([0, 1, 0, 1])
// weights = torch.tensor([3., 7.]).gather(0, target)
// loss = nn.BCELoss(weights)
// sigmoid = nn.Sigmoid()
// out = loss(sigmoid(input), target.float()) # tensor(3.1531)
    fn test_binary_cross_entropy_with_smoothing() {
// import torch
// from torch import nn
// input = torch.tensor([0.8271, 0.9626, 0.3796, 0.2355])
// target = torch.tensor([0., 1., 0., 1.])
// target_smooth = target * (1 - 0.1) + (0.1 / 2)
// loss = nn.BCELoss()
// sigmoid = nn.Sigmoid()
// out = loss(sigmoid(input), target_smooth) # tensor(0.7490)
    fn test_binary_cross_entropy_multilabel() {
// import torch
// from torch import nn
// input = torch.tensor([[0.5150, 0.3097, 0.7556], [0.4974, 0.9879, 0.1564]])
// target = torch.tensor([[1., 0., 1.], [1., 0., 0.]])
// weights = torch.tensor([3., 7., 0.9])
// loss = nn.BCEWithLogitsLoss()
// out = loss(input, target) # tensor(0.7112)
    fn test_binary_cross_entropy_multilabel_with_weights() {
// import torch
// from torch import nn
// input = torch.tensor([[0.5150, 0.3097, 0.7556], [0.4974, 0.9879, 0.1564]])
// target = torch.tensor([[1., 0., 1.], [1., 0., 0.]])
// loss = nn.BCEWithLogitsLoss()
// out = loss(input, target) # tensor(3.1708)
    fn test_binary_cross_entropy_multilabel_with_smoothing() {
// import torch
// from torch import nn
// input = torch.tensor([[0.5150, 0.3097, 0.7556], [0.4974, 0.9879, 0.1564]])
// target = torch.tensor([[1., 0., 1.], [1., 0., 0.]])
// target_smooth = target * (1 - 0.1) + (0.1 / 3)
// loss = nn.BCELoss()
// sigmoid = nn.Sigmoid()
// out = loss(sigmoid(input), target_smooth) # tensor(0.7228)
    fn multilabel_weights_should_match_target() {
// import torch
// from torch import nn
// input = torch.tensor([[0.5150, 0.3097, 0.7556], [0.4974, 0.9879, 0.1564]])
// target = torch.tensor([[1., 0., 1.], [1., 0., 0.]])
// loss = nn.BCEWithLogitsLoss()
// out = loss(input, target) # tensor(3.1708)
    fn display() {
/// Configuration for CosineEmbeddingLoss.
pub struct CosineEmbeddingLossConfig {
/// Margin for negative samples.
/// Specifies the reduction to apply to the output.
impl CosineEmbeddingLossConfig {
/// Initialize CosineEmbeddingLoss.
    pub fn init(&self) -> CosineEmbeddingLoss {
/// Cosine embedding loss between two tensors.
///
/// Measures cosine distance between tensors.
/// Used for learning embeddings or similarity.
pub struct CosineEmbeddingLoss {
/// Margin value. Default: 0.0
/// Reduction method
impl Default for CosineEmbeddingLoss {
    fn default() -> Self {
impl ModuleDisplay for CosineEmbeddingLoss {
    fn custom_settings(&self) -> Option<DisplaySettings> {
    fn custom_content(&self, content: Content) -> Option<Content> {
impl CosineEmbeddingLoss {
/// Creates a new instance
    pub fn new() -> Self {
/// Compute loss with reduction.
///
/// # Shapes
///
/// - input1: ``[batch_size, embedding_dim]``
/// - input2: ``[batch_size, embedding_dim]``
/// - target: ``[batch_size]`` with values 1 or -1
///
/// # Returns
///
/// Loss tensor of shape ``[1]``
    pub fn forward<B: Backend>(
        &self,
        input1: Tensor<B, 2>,
        input2: Tensor<B, 2>,
        target: Tensor<B, 1, Int>,
    ) -> Tensor<B, 1> {
/// Compute loss without applying reduction.
///
/// # Arguments
///
/// * `input1` - First input tensor of shape ``[batch_size, embedding_dim]``
/// * `input2` - Second input tensor of shape ``[batch_size, embedding_dim]``
/// * `target` - Target tensor of shape ``[batch_size]`` with values 1 or -1
///
/// # Returns
///
/// Tensor of per-element losses with shape ``[batch_size]``
    pub fn forward_no_reduction<B: Backend>(
        &self,
        input1: Tensor<B, 2>,
        input2: Tensor<B, 2>,
        target: Tensor<B, 1, Int>,
    ) -> Tensor<B, 1> {
// cos_sim shape: [batch_size, 1]
// cos_sim shape: [batch_size]
// Similar pairs (target == 1) - Formula: L = 1 - cos_sim
// Dissimilar pairs (target == -1) - Formula: L = max(0, cos_sim - margin)
// return loss shape: [batch_size]
    fn assertions<B: Backend>(
        &self,
        input1: &Tensor<B, 2>,
        input2: &Tensor<B, 2>,
        target: &Tensor<B, 1, Int>,
    ) {
    type FT = FloatElem<TestBackend>;
    fn cosine_embedding_loss_positive_target() {
// Two identical vectors should have cosine similarity of 1
// Target 1 means that inputs should be similar
// For identical vectors, 1 - cos_sim = 1 - 1 = 0
    fn cosine_embedding_loss_negative_target() {
// Two identical vectors should have cosine similarity of 1
// Target -1 means that inputs should be dissimilar
// With margin 0.0, max(0, cos_sim - margin) = max(0, 1 - 0) = 1
// Create a loss with Sum reduction for testing
// With margin 0.5, max(0, cos_sim - margin) = max(0, 1 - 0.5) = 0.5
    fn cosine_embedding_loss_mixed_targets() {
// Mixed targets
    fn display() {
/// Configuration to create a [Cross-entropy loss](CrossEntropyLoss) using the [init function](CrossEntropyLossConfig::init).
pub struct CrossEntropyLossConfig {
/// Create padded cross entropy.
///
/// Prevents pad tokens from impacting loss calculation.
/// Create weighted cross-entropy.
///
/// The loss of a specific sample will simply be given by: weight * log(p(x)) * 1,
///
/// # Pre-conditions
///   - The order of the weight vector should correspond to the label integer assignment.
///   - Targets assigned negative Int's will not be allowed.
/// Create cross-entropy with label smoothing.
///
/// Hard labels {0, 1} will be changed to y_smoothed = y(1 - a) + a / nr_classes.
/// Alpha = 0 would be the same as default.
/// Create cross-entropy with probabilities as input instead of logits.
///
impl CrossEntropyLossConfig {
/// Initialize [Cross-entropy loss](CrossEntropyLoss).
    pub fn init<B: Backend>(&self, device: &B::Device) -> CrossEntropyLoss<B> {
    fn assertions(&self) {
/// Calculate the cross entropy loss from the input logits and the targets.
///
/// Should be created using [CrossEntropyLossConfig]
pub struct CrossEntropyLoss<B: Backend> {
/// Pad tokens to ignore in the loss calculation.
/// Weights for cross-entropy.
/// Label smoothing factor.
/// Use logits as input.
impl<B: Backend> ModuleDisplay for CrossEntropyLoss<B> {
    fn custom_settings(&self) -> Option<DisplaySettings> {
    fn custom_content(&self, content: Content) -> Option<Content> {
impl<B: Backend> CrossEntropyLoss<B> {
/// For backward compatibility.
    pub fn new(pad_index: Option<usize>, device: &B::Device) -> Self {
/// Compute the criterion on the input tensor.
///
/// # Shapes
///
/// - logits: `[batch_size, num_targets]`
/// - targets: `[batch_size]`
    pub fn forward(&self, logits: Tensor<B, 2>, targets: Tensor<B, 1, Int>) -> Tensor<B, 1> {
    fn forward_smoothed(
        &self,
        logits: Tensor<B, 2>,
        targets: Tensor<B, 1, Int>,
        alpha: f32,
    ) -> Tensor<B, 1> {
    fn forward_default(&self, logits: Tensor<B, 2>, targets: Tensor<B, 1, Int>) -> Tensor<B, 1> {
    fn compute_smoothed_targets(
        shape: [usize;
    fn padding_mask(&self, targets: &Tensor<B, 1, Int>) -> Option<Tensor<B, 1, Bool>> {
    fn apply_mask_1d(mut tensor: Tensor<B, 1>, mask: Option<Tensor<B, 1, Bool>>) -> Tensor<B, 1> {
    fn apply_mask_2d(mut tensor: Tensor<B, 2>, mask: Option<Tensor<B, 1, Bool>>) -> Tensor<B, 2> {
    fn assertions(logits: Tensor<B, 2>, targets: Tensor<B, 1, Int>) {
    type FT = FloatElem<TestBackend>;
    fn test_cross_entropy_loss_with_weights() {
    fn test_label_smoothing_with_weights_and_alpha_zero() {
    fn test_cross_entropy_loss() {
    fn test_label_smoothing_alpha_equal_zero() {
    fn test_cross_entropy_loss_with_pad_token() {
    fn test_label_smoothing_with_zero_alpha_and_pad_token() {
    fn test_label_smoothing_target_conversion() {
    fn test_label_smoothing() {
    fn display() {
/// Configuration to create a [Huber loss](HuberLoss).
pub struct HuberLossConfig {
/// The bound where the Huber loss function changes from quadratic to linear behaviour.
impl HuberLossConfig {
/// Initialize [Huber loss](HuberLoss).
    pub fn init(&self) -> HuberLoss {
    fn assertions(&self) {
// This also tests for normality
/// Calculate the Huber loss between the inputs and the target.
///
/// The loss for each element of the residuals `r = targets - predictions` is given by
///
/// ```text
/// L(r) = 0.5 * r^2                  if |r| <= d
/// L(r) = 0.5 * d^2 + d * (|r| - d)  if |r| >  d
/// ```
///
/// where `d` is the configured `delta`. In particular, this is equal to the
/// [L2 Loss](super::MseLoss) for residuals with magnitude smaller than `delta`,
/// but behaves linearly instead of quadratically for large residuals.
///
/// This loss function is less sensitive to outliers than the mean squared error loss.
///
/// See also: <https://en.wikipedia.org/wiki/Huber_loss>
pub struct HuberLoss {
/// The bound where the Huber loss function changes from quadratic to linear behaviour.
/// Precomputed value for the linear bias.
// delta * delta * 0.5 precomputed
impl ModuleDisplay for HuberLoss {
    fn custom_settings(&self) -> Option<DisplaySettings> {
    fn custom_content(&self, content: Content) -> Option<Content> {
impl HuberLoss {
/// Compute the loss element-wise for the predictions and targets, then reduce
/// to a single loss value.
///
/// `Reduction::Auto` behaves as `Reduction::Mean`.
///
/// # Shapes
///
/// - predictions: \[...dims\]
/// - targets: \[...dims\]
/// - output: \[1\]
    pub fn forward<const D: usize, B: Backend>(
        &self,
        predictions: Tensor<B, D>,
        targets: Tensor<B, D>,
        reduction: Reduction,
    ) -> Tensor<B, 1> {
/// Compute the loss element-wise for the predictions and targets.
///
/// # Shapes
///
/// - predictions: [...dims]
/// - targets: [...dims]
/// - output: [...dims]
    pub fn forward_no_reduction<const D: usize, B: Backend>(
        &self,
        predictions: Tensor<B, D>,
        targets: Tensor<B, D>,
    ) -> Tensor<B, D> {
/// Compute the loss element-wise for the given residuals.
///
/// # Shapes
///
/// - residuals: [...dims]
/// - output: [...dims]
    pub fn forward_residuals<const D: usize, B: Backend>(
        &self,
        residuals: Tensor<B, D>,
    ) -> Tensor<B, D> {
// We are interested in `sign(r)` when `abs(r) > self.delta`. Note that the
// `sign()` function, in general, suffers from a jump at 0.
// Instead the following tensor implements `delta * sign(r)` for values outside
// the bound:
// 0.5 * d^2 + d * (|r| - d) =
// d * |r| - 0.5 * d^2
// Moreover |r| = sign(r) * r
    type TestTensor<const D: usize> = Tensor<TestBackend, D>;
    type FT = FloatElem<TestBackend>;
    fn test_huber_loss() {
    fn test_huber_ad_loss() {
        type TestAutodiffTensor = Tensor<crate::TestAutodiffBackend, 1>;
    fn display() {
/// Calculate the mean squared error loss from the input logits and the targets.
pub struct MseLoss;
impl Default for MseLoss {
    fn default() -> Self {
impl MseLoss {
/// Create the criterion.
    pub fn new() -> Self {
/// Compute the criterion on the input tensor.
///
/// # Shapes
///
/// - logits: [batch_size, num_targets]
/// - targets: [batch_size, num_targets]
    pub fn forward<const D: usize, B: Backend>(
        &self,
        logits: Tensor<B, D>,
        targets: Tensor<B, D>,
        reduction: Reduction,
    ) -> Tensor<B, 1> {
/// Compute the criterion on the input tensor without reducing.
    pub fn forward_no_reduction<const D: usize, B: Backend>(
        &self,
        logits: Tensor<B, D>,
        targets: Tensor<B, D>,
    ) -> Tensor<B, D> {
    fn test_mse_loss() {
    fn display() {
/// Configuration for creating a [PoissonNllLoss](PoissonNllLoss) instance.
///
/// This configuration allows customization of the Poisson Negative Log Likelihood (NLL) loss
/// behavior, such as whether the input is in log-space, whether to include the Stirling
/// approximation term, and a small epsilon value to avoid numerical instability.
pub struct PoissonNllLossConfig {
/// If `true`, the predictions are expected to be in log-space.
///
/// When `log_input` is `true`, the loss is computed as:
/// ```text
/// L(predictions, target) = exp(predictions) - target * predictions
/// ```
/// When `log_input` is `false`, the loss is computed as:
/// ```text
/// L(predictions, target) = predictions - target * log(predictions + eps)
/// ```
/// Whether to compute the full loss, including the Stirling approximation term.
///
/// When `full` is `true`, the Stirling approximation term is added to the loss:
/// ```text
/// target * log(target) - target + 0.5 * log(2 * PI * target)
/// ```
/// A small value to avoid evaluation of `log(0)` when `log_input` is `false`.
///
/// This epsilon value is added to the predictions to ensure numerical stability
/// when computing the logarithm.
impl PoissonNllLossConfig {
/// Initializes a [PoissonNllLoss](PoissonNllLoss) instance with the current configuration.
///
/// # Panics
/// - Panics if `eps` is not a positive number.
    pub fn init(&self) -> PoissonNllLoss {
/// Validates the configuration parameters.
///
/// # Panics
/// - Panics if `eps` is not a positive number.
    fn assertions(&self) {
/// Negative Log Likelihood (NLL) loss with a Poisson distribution assumption for the target.
///
/// This loss function is used when the target values are assumed to follow a Poisson distribution.
/// The loss is defined as:
/// ```text
/// target ~ Poisson(input)
/// L(predictions, target) = predictions - target * log(predictions) + log(target!)
/// ```
/// The last term (`log(target!)`) can be omitted or approximated using Stirling's formula.
/// The approximation is applied for `target > 1`, while for `target <= 1`, zeros are added to the loss.
///
/// For more details, see:
/// <https://en.wikipedia.org/wiki/Poisson_regression#Maximum_likelihood-based_parameter_estimation>
pub struct PoissonNllLoss {
/// If `true`, the predictions are expected to be in log-space.
/// Whether to compute the full loss, including the Stirling approximation term.
/// A small value to avoid evaluation of `log(0)` when `log_input` is `false`.
impl ModuleDisplay for PoissonNllLoss {
    fn custom_settings(&self) -> Option<DisplaySettings> {
    fn custom_content(&self, content: Content) -> Option<Content> {
impl PoissonNllLoss {
/// Computes the loss element-wise for the given predictions and targets, then reduces
/// the result to a single loss value.
///
/// # Arguments
/// - `predictions`: The predicted values.
/// - `targets`: The target values.
/// - `reduction`: The reduction method to apply. `Reduction::Auto` behaves as `Reduction::Mean`.
///
/// # Shapes
/// - `predictions`: `[...dims]`
/// - `targets`: `[...dims]`
/// - `output`: `[1]`
///
/// # Panics
/// - Panics if the shapes of `predictions` and `targets` do not match.
/// - Panics if any target value is negative.
/// - Panics if `log_input` is `false` and any prediction value is negative.
    pub fn forward<const D: usize, B: Backend>(
        &self,
        predictions: Tensor<B, D>,
        targets: Tensor<B, D>,
        reduction: Reduction,
    ) -> Tensor<B, 1> {
/// Computes the loss element-wise for the given predictions and targets without reduction.
///
/// # Arguments
/// - `predictions`: The predicted values.
/// - `targets`: The target values.
///
/// # Shapes
/// - `predictions`: `[...dims]`
/// - `targets`: `[...dims]`
/// - `output`: `[...dims]`
///
/// # Panics
/// - Panics if the shapes of `predictions` and `targets` do not match.
/// - Panics if any target value is negative.
/// - Panics if `log_input` is `false` and any prediction value is negative.
    pub fn forward_no_reduction<const D: usize, B: Backend>(
        &self,
        predictions: Tensor<B, D>,
        targets: Tensor<B, D>,
    ) -> Tensor<B, D> {
/// Validates the input tensors for the loss computation.
///
/// # Panics
/// - Panics if the shapes of `predictions` and `targets` do not match.
/// - Panics if any target value is negative.
/// - Panics if `log_input` is `false` and any prediction value is negative.
    fn assertions<const D: usize, B: Backend>(
        &self,
        predictions: &Tensor<B, D>,
        targets: &Tensor<B, D>,
    ) {
    type TestTensor<const D: usize> = Tensor<TestBackend, D>;
    type FT = FloatElem<TestBackend>;
    fn test_poisson_nll_loss() {
    fn test_poisson_nll_loss_no_log_input() {
    fn test_poisson_nll_loss_full() {
    fn test_poisson_nll_loss_gradients() {
        type TestAutodiffTensor = Tensor<crate::TestAutodiffBackend, 1>;
    fn test_negative_eps() {
    fn test_targets_with_negative_values() {
    fn test_shape_tensors() {
    fn test_exp_predictions_non_negative() {
    fn display() {
/// The reduction type for the loss.
pub enum Reduction {
/// The mean of the losses will be returned.
/// The sum of the losses will be returned.
/// The mean of the losses will be returned.
/// Generate an autoregressive attention mask.
///
/// The mask can be used in Transformer modules to train models to generate tensors sequentially.
pub fn generate_autoregressive_mask<B: Backend>(
    batch_size: usize,
    seq_length: usize,
    device: &B::Device,
) -> Tensor<B, 3, Bool> {
/// Generate a padding attention mask.
pub struct GeneratePaddingMask<B: Backend> {
/// The generated tensor.
/// The generated mask.
/// Defines an enumeration to specify sequence length options for padding
pub enum SeqLengthOption {
/// No maximum length; use the longest sequence
/// Maximum length specified, truncate if necessary
/// Fixed length, pad or truncate to this exact length
impl From<Option<usize>> for SeqLengthOption {
    fn from(val: Option<usize>) -> Self {
/// Generates a padding attention mask for a batch of token sequences.
///
/// # Arguments
///
/// * `pad_token` - The token ID used for padding
/// * `tokens_list` - Vector of token sequences (each sequence is a vector of token IDs)
/// * `seq_length` - Sequence length option (NoMax, Max, or Fixed)
/// * `device` - The device for tensor operations
///
/// # Returns
///
/// A `GeneratePaddingMask` containing the padded tensor and corresponding mask
pub fn generate_padding_mask<B: Backend>(
    pad_token: usize,
    tokens_list: Vec<Vec<usize>>,
    seq_length: impl Into<SeqLengthOption>,
    device: &B::Device,
) -> GeneratePaddingMask<B> {
    fn test_generate_autoregressive_mask() {
    fn test_generate_padding_mask() {
/// Configuration to create a [Multi Head Attention](MultiHeadAttention) layer using the [init function](MultiHeadAttentionConfig::init).
pub struct MultiHeadAttentionConfig {
/// The size of each linear layer.
/// The number of heads.
/// The dropout rate. Default: 0.1
/// The minimum value a float can take. Default: -1.0e4
/// This is used to mask attention scores before calculating attention weights.
/// A value too low might result in NaN.
/// Use "quiet softmax" instead of regular softmax.
///
/// - Usage may improve performance by allowing attention heads to deposit no information (if the sequence contains no information relevant to that head).
/// - Usage may reduce the entropy of weights in the model, enhancing quantization and compression.
///
/// Reference: <https://www.evanmiller.org/attention-is-off-by-one.html>
/// The type of function used to initialize neural network parameters
/// The multihead attention module as describe in the paper [Attention Is All You Need](https://arxiv.org/abs/1706.03762).
///
/// # Params
///
/// - `query`: [`Linear`] layer with `d_model` input and output features.
/// - `key`: [`Linear`] layer with `d_model` input and output features.
/// - `value`: [`Linear`] layer with `d_model` input and output features.
/// - `output`: [`Linear`] layer with `d_model` input and output features.
///
/// Should be created with [MultiHeadAttentionConfig].
pub struct MultiHeadAttention<B: Backend> {
/// Linear layer to transform the input features into the query space.
/// Linear layer to transform the input features into the key space.
/// Linear layer to transform the input features into the value space.
/// Linear layer to transform the output features back to the original space.
/// Dropout layer.
/// Activation function.
/// The size of each linear layer.
/// The number of heads.
/// Size of the key and query vectors.
/// Minimum value a float can take.
/// Use "quiet softmax" instead of regular softmax.
impl<B: Backend> ModuleDisplay for MultiHeadAttention<B> {
    fn custom_settings(&self) -> Option<DisplaySettings> {
    fn custom_content(&self, content: Content) -> Option<Content> {
/// [Multihead attention](MultiHeadAttention) forward pass input argument.
pub struct MhaInput<B: Backend> {
/// Shape `[batch_size, seq_length_1, d_model]`
/// Shape `[batch_size, seq_length_2, d_model]`
/// Shape `[batch_size, seq_length_2, d_model]`
impl MultiHeadAttentionConfig {
/// Initialize a new [multihead attention](MultiHeadAttention) module.
    pub fn init<B: Backend>(&self, device: &B::Device) -> MultiHeadAttention<B> {
impl<B: Backend> MhaInput<B> {
/// Create a [multihead attention](MultiHeadAttention) input argument
/// by setting the query, key and value to the given tensor.
///
/// # Shape
/// - tensor: `[batch_size, seq_length, d_model]`
    pub fn self_attn(tensor: Tensor<B, 3>) -> Self {
/// Create a [multihead attention](MultiHeadAttention) input argument.
    pub fn new(query: Tensor<B, 3>, key: Tensor<B, 3>, value: Tensor<B, 3>) -> Self {
/// Register the padding mask.
    pub fn mask_pad(mut self, mask_pad: Tensor<B, 2, Bool>) -> Self {
/// Register the attention mask.
    pub fn mask_attn(mut self, mask_attn: Tensor<B, 3, Bool>) -> Self {
/// [Multihead attention](MultiHeadAttention) outputs.
pub struct MhaOutput<B: Backend> {
/// The attention weights `[batch_size, n_heads, seq_length_1, seq_length_2]`.
/// The context tensor `[batch_size, seq_length_1, d_model]`.
impl<B: Backend> MultiHeadAttention<B> {
/// Applies the forward pass on the input tensors.
///
/// See [MultiHeadAttention](MultiHeadAttention) for more information.
///
/// # Shapes
///
/// - query: `[batch_size, seq_length_1, d_model]`
/// - key: `[batch_size, seq_length_2, d_model]`
/// - value: `[batch_size, seq_length_2, d_model]`
/// - output: `[batch_size, seq_length_1, d_model]`
    pub fn forward(&self, input: MhaInput<B>) -> MhaOutput<B> {
/// Applies the forward pass using a cache.
///
/// # Shapes
///
/// - query: `[batch_size, seq_length_1, d_model]`
/// - key: `[batch_size, seq_length_2, d_model]`
/// - value: `[batch_size, seq_length_2, d_model]`
/// - output: `[batch_size, seq_length_1, d_model]`
    pub fn forward_cache(&self, input: MhaInput<B>, cache: &mut MhaCache<B>) -> MhaOutput<B> {
    fn attn_scores(&self, query: Tensor<B, 4>, key: Tensor<B, 4>) -> Tensor<B, 4> {
    fn attn_weights(
        &self,
        mut attn_scores: Tensor<B, 4>,
        mask_pad: Option<Tensor<B, 2, Bool>>,
        mask_attn: Option<Tensor<B, 3, Bool>>,
    ) -> Tensor<B, 4> {
    fn attention_linear(&self, x: Tensor<B, 3>, linear: &Linear<B>) -> Tensor<B, 4> {
/// Cache for the [Multi Head Attention](MultiHeadAttention) layer.
///
/// To be used during inference when decoding tokens.
pub struct MhaCache<B: Backend> {
enum MhaLinearCache<B: Backend, const D: usize> {
impl<B: Backend> MhaCache<B> {
/// Initialize a cache for autoregressive inference.
    pub fn autoregressive() -> Self {
/// Initialize a cache for autoregressive inference, but with a fixed memory used for keys and
/// values (cross-attention).
    pub fn autoregressive_cross_attention() -> Self {
impl<B: Backend, const D: usize> MhaLinearCache<B, D> {
    pub fn forward<F: Fn(Tensor<B, 3>) -> Tensor<B, D>>(
        &mut self,
        tensor: Tensor<B, 3>,
        func: F,
    ) -> Tensor<B, D> {
    fn test_self_attention_shapes() {
    fn test_generic_mha_shapes() {
    fn test_self_attention_mask_pad() {
// Create a padding mask
// Change the end of the tensor
// Check that the beginning of each tensor is the same
    fn test_autoregressive_mask_should_have_same_output_as_autoregressive_decoding() {
    fn display() {
impl<B: Backend, const D: usize> TensorCache<B, D> {
    pub(crate) fn forward_autoregressive<F>(
        &mut self,
        tensor: Tensor<B, 3>,
        dim_cat: usize,
        func: F,
    ) -> Tensor<B, D>
    where
        F: Fn(Tensor<B, 3>) -> Tensor<B, D>,
    {
    pub(crate) fn forward_full<F>(&mut self, tensor: Tensor<B, 3>, func: F) -> Tensor<B, D>
    where
        F: Fn(Tensor<B, 3>) -> Tensor<B, D>,
    {
pub(crate) enum CacheState<T> {
/// A cache for a tensor.
pub struct TensorCache<B: Backend, const D: usize> {
impl<B: Backend, const D: usize> TensorCache<B, D> {
/// Creates a new empty cache.
///
/// # Returns
///
/// The empty cache.
    pub fn empty() -> Self {
pub(crate) fn checks_channels_div_groups(channels_in: usize, channels_out: usize, groups: usize) {
// https://github.com/tracel-ai/burn/issues/2676
/// Only symmetric padding is currently supported. As such, using `Same` padding with an even kernel
/// size is not supported as it will not produce the same output size.
pub(crate) fn check_same_padding_support(kernel_size: &[usize]) {
/// Configuration to create a [1D convolution](Conv1d) layer using the [init function](Conv1dConfig::init).
pub struct Conv1dConfig {
/// The number of input channels.
/// The number of output channels.
/// The size of the kernel.
/// The stride of the convolution.
/// Spacing between kernel elements.
/// Controls the connections between input and output channels.
/// The padding configuration.
///
/// ### Warning
/// Only symmetric padding is currently supported. As such, using `Same` padding with an even kernel
/// size is not supported as it will not produce the same output size.
/// If bias should be added to the output.
/// The type of function used to initialize neural network parameters
/// Applies a 1D convolution over input tensors.
///
/// Should be created with [Conv1dConfig].
pub struct Conv1d<B: Backend> {
/// Tensor of shape `[channels_out, channels_in / groups, kernel_size]`
/// Tensor of shape `[channels_out]`
/// Stride of the convolution.
/// Size of the kernel.
/// Spacing between kernel elements.
/// Controls the connections between input and output channels.
/// Padding configuration.
impl<B: Backend> ModuleDisplay for Conv1d<B> {
    fn custom_settings(&self) -> Option<DisplaySettings> {
    fn custom_content(&self, content: Content) -> Option<Content> {
// Format padding
// Format stride/dilation as strings
// Extract channels in/out from weight dims
impl Conv1dConfig {
/// Initialize a new [conv1d](Conv1d) module.
    pub fn init<B: Backend>(&self, device: &B::Device) -> Conv1d<B> {
impl<B: Backend> Conv1d<B> {
/// Applies the forward pass on the input tensor.
///
/// See [conv1d](burn::tensor::module::conv1d) for more information.
///
/// # Shapes
///
/// - input: `[batch_size, channels_in, length_in]`
/// - output: `[batch_size, channels_out, length_out]`
    pub fn forward(&self, input: Tensor<B, 3>) -> Tensor<B, 3> {
    type FT = FloatElem<TestBackend>;
    fn initializer_default() {
    fn initializer_zeros() {
    fn same_with_even_kernel_is_invalid() {
    fn display() {
    fn input_channels_mismatch() {
/// Configuration to create a [2D convolution](Conv2d) layer, using the [init function](Conv2dConfig::init).
pub struct Conv2dConfig {
/// The number of channels.
/// The size of the kernel.
/// The stride of the convolution.
/// Spacing between kernel elements.
/// Controls the connections between input and output channels.
/// The padding configuration.
///
/// ### Warning
/// Only symmetric padding is currently supported. As such, using `Same` padding with an even kernel
/// size is not supported as it will not produce the same output size.
/// If bias should be added to the output.
/// The type of function used to initialize neural network parameters
/// Applies a 2D convolution over input tensors.
///
/// Should be created with [Conv2dConfig].
pub struct Conv2d<B: Backend> {
/// Tensor of shape `[channels_out, channels_in / groups, kernel_size_1, kernel_size_2]`
/// Tensor of shape `[channels_out]`
/// Stride of the convolution.
/// Size of the kernel.
/// Spacing between kernel elements.
/// Controls the connections between input and output channels.
/// The padding configuration.
impl Conv2dConfig {
/// Initialize a new [conv2d](Conv2d) module.
    pub fn init<B: Backend>(&self, device: &B::Device) -> Conv2d<B> {
impl<B: Backend> ModuleDisplay for Conv2d<B> {
    fn custom_settings(&self) -> Option<DisplaySettings> {
    fn custom_content(&self, content: Content) -> Option<Content> {
// Since padding does not implement ModuleDisplay, we need to format it manually.
// Format the stride, kernel_size and dilation as strings, formatted as arrays instead of indexed.
impl<B: Backend> Conv2d<B> {
/// Applies the forward pass on the input tensor.
///
/// See [conv2d](burn::tensor::module::conv2d) for more information.
///
/// # Shapes
/// - `input`: `[batch_size, channels_in, height_in, width_in]`
/// - `output`: `[batch_size, channels_out, height_out, width_out]`
///
/// # Example
/// ```rust,ignore
/// use burn::nn::conv::Conv2dConfig;
/// use burn::tensor::Tensor;
///
/// // Assuming backend type alias `B`
/// let device = Default::default();
/// let conv = Conv2dConfig::new([3, 8], [3, 3]).init::<B>(&device);
///
/// let x = Tensor::<B, 4>::zeros([1, 3, 28, 28], &device);
/// let y = conv.forward(x);
///
/// println!("{:?}", y.dims()); // [1, 8, 26, 26]
/// ```
    pub fn forward(&self, input: Tensor<B, 4>) -> Tensor<B, 4> {
    type FT = FloatElem<TestBackend>;
    fn initializer_default() {
    fn initializer_zeros() {
    fn initializer_fan_out() {
// test that fan_out is passed to `init_with()`
    fn initializer_fan_with_groups_is_valid() {
    fn channels_with_groups_is_invalid() {
    fn same_with_even_kernel_is_invalid() {
    fn display() {
    fn input_channels_mismatch() {
/// Configuration to create a [3D convolution](Conv3d) layer, using the [init function](Conv3dConfig::init).
pub struct Conv3dConfig {
/// The number of channels.
/// The size of the kernel.
/// The stride of the convolution.
/// Spacing between kernel elements.
/// Controls the connections between input and output channels.
/// The padding configuration.
/// If bias should be added to the output.
/// The type of function used to initialize neural network parameters
/// Applies a 3D convolution over input tensors.
///
/// Should be created with [Conv3dConfig].
pub struct Conv3d<B: Backend> {
/// Tensor of shape `[channels_out, channels_in / groups, kernel_size_1, kernel_size_2, kernel_size_3]`
/// Tensor of shape `[channels_out]`
/// Stride of the convolution.
/// Size of the kernel.
/// Spacing between kernel elements.
/// Controls the connections between input and output channels.
/// The padding configuration.
impl Conv3dConfig {
/// Initialize a new [conv3d](Conv3d) module.
    pub fn init<B: Backend>(&self, device: &B::Device) -> Conv3d<B> {
impl<B: Backend> ModuleDisplay for Conv3d<B> {
    fn custom_settings(&self) -> Option<DisplaySettings> {
    fn custom_content(&self, content: Content) -> Option<Content> {
// Padding doesn't implement ModuleDisplay, so format manually.
// Format arrays as strings (consistent with Conv2d/Conv1d).
// Weight dims: [channels_out, channels_in/groups, k1, k2, k3]
impl<B: Backend> Conv3d<B> {
/// Applies the forward pass on the input tensor.
///
/// See [conv3d](burn::tensor::module::conv3d) for more information.
///
/// # Shapes
///
/// - input: `[batch_size, channels_in, depth_in, height_in, width_in]`
/// - output: `[batch_size, channels_out, depth_out, height_out, width_out]`
    pub fn forward(&self, input: Tensor<B, 5>) -> Tensor<B, 5> {
    type FT = FloatElem<TestBackend>;
    fn initializer_default() {
    fn initializer_zeros() {
    fn initializer_fan_out() {
// test that fan_out is passed to `init_with()`
    fn initializer_fan_with_groups_is_valid() {
    fn same_with_even_kernel_is_invalid() {
    fn display() {
    fn input_channels_mismatch() {
/// Configuration to create an [1D transposed convolution](ConvTranspose1d) layer
/// using the [init function](ConvTranspose1dConfig::init).
pub struct ConvTranspose1dConfig {
/// The number of channels.
/// The size of the kernel.
/// The stride of the convolution.
/// Spacing between kernel elements.
/// Controls the connections between input and output channels.
/// The padding configuration.
/// The padding output configuration.
/// If bias should be added to the output.
/// The type of function used to initialize neural network parameters
/// Applies a 1D transposed convolution over input tensors.
pub struct ConvTranspose1d<B: Backend> {
/// Tensor of shape `[channels_in, channels_out / groups, kernel_size]`
/// Tensor of shape `[channels_out]`
/// Stride of the convolution.
/// Size of the kernel.
/// Spacing between kernel elements.
/// Controls the connections between input and output channels.
/// The padding configuration.
/// The padding output configuration.
/// The number of channels.
impl<B: Backend> ModuleDisplay for ConvTranspose1d<B> {
    fn custom_settings(&self) -> Option<DisplaySettings> {
    fn custom_content(&self, content: Content) -> Option<Content> {
impl ConvTranspose1dConfig {
/// Initialize a new [conv transpose 1d](ConvTranspose1d) module.
    pub fn init<B: Backend>(&self, device: &B::Device) -> ConvTranspose1d<B> {
impl<B: Backend> ConvTranspose1d<B> {
/// Applies the forward pass on the input tensor.
///
/// See also [conv_transpose1d](burn::tensor::module::conv_transpose1d).
///
/// # Shapes
///
/// - input: `[batch_size, channels_in, length_in]`
/// - output: `[batch_size, channels_out, length_out]`
    pub fn forward(&self, input: Tensor<B, 3>) -> Tensor<B, 3> {
    type FT = FloatElem<TestBackend>;
    fn initializer_default() {
    fn initializer_zeros() {
    fn display() {
    fn input_channels_mismatch() {
/// Configuration to create an [2D transposed convolution](ConvTranspose2d) layer
/// using the [init function](ConvTranspose2dConfig::init).
pub struct ConvTranspose2dConfig {
/// The number of channels.
/// The size of the kernel.
/// The stride of the convolution.
/// Spacing between kernel elements.
/// Controls the connections between input and output channels.
/// The padding configuration.
/// The padding output configuration.
/// If bias should be added to the output.
/// The type of function used to initialize neural network parameters
/// Applies a 2D transposed convolution over input tensors.
pub struct ConvTranspose2d<B: Backend> {
/// Tensor of shape `[channels_in, channels_out / groups, kernel_size_1, kernel_size_2]`
/// Tensor of shape `[channels_out]`
/// Stride of the convolution.
/// Size of the kernel.
/// Spacing between kernel elements.
/// Controls the connections between input and output channels.
/// Padding configuration.
/// Padding output configuration.
/// Number of channels.
impl<B: Backend> ModuleDisplay for ConvTranspose2d<B> {
    fn custom_settings(&self) -> Option<DisplaySettings> {
    fn custom_content(&self, content: Content) -> Option<Content> {
impl ConvTranspose2dConfig {
/// Initialize a new [conv transpose 2d](ConvTranspose2d) module.
    pub fn init<B: Backend>(&self, device: &B::Device) -> ConvTranspose2d<B> {
impl<B: Backend> ConvTranspose2d<B> {
/// Applies the forward pass on the input tensor.
///
/// See also [conv_transpose2d](burn::tensor::module::conv_transpose2d).
///
/// # Shapes
///
/// - input: `[batch_size, channels_in, height_in, width_in]`
/// - output: `[batch_size, channels_out, height_out, width_out]`
    pub fn forward(&self, input: Tensor<B, 4>) -> Tensor<B, 4> {
    type FT = FloatElem<TestBackend>;
    fn initializer_default() {
    fn initializer_zeros() {
    fn display() {
    fn input_channels_mismatch() {
/// Configuration to create an [3D transposed convolution](ConvTranspose3d) layer
/// using the [init function](ConvTranspose3dConfig::init).
pub struct ConvTranspose3dConfig {
/// The number of channels.
/// The size of the kernel.
/// The stride of the convolution.
/// Spacing between kernel elements.
/// Controls the connections between input and output channels.
/// The padding configuration.
/// The padding output configuration.
/// If bias should be added to the output.
/// The type of function used to initialize neural network parameters
/// Applies a 3D transposed convolution over input tensors.
pub struct ConvTranspose3d<B: Backend> {
/// Tensor of shape `[channels_in, channels_out / groups, kernel_size_1, kernel_size_2, kernel_size_3]`
/// Tensor of shape `[channels_out]`
/// Stride of the convolution.
/// Size of the kernel.
/// Spacing between kernel elements.
/// Controls the connections between input and output channels.
/// Padding configuration.
/// Padding output configuration.
/// Number of channels.
impl<B: Backend> ModuleDisplay for ConvTranspose3d<B> {
    fn custom_settings(&self) -> Option<DisplaySettings> {
    fn custom_content(&self, content: Content) -> Option<Content> {
impl ConvTranspose3dConfig {
/// Initialize a new [conv transpose 2d](ConvTranspose3d) module.
    pub fn init<B: Backend>(&self, device: &B::Device) -> ConvTranspose3d<B> {
impl<B: Backend> ConvTranspose3d<B> {
/// Applies the forward pass on the input tensor.
///
/// See also [conv_transpose3d](burn::tensor::module::conv_transpose3d).
///
/// # Shapes
///
/// - input: `[batch_size, channels_in, depth_in, height_in, width_in]`
/// - output: `[batch_size, channels_out, depth_out, height_out, width_out]`
    pub fn forward(&self, input: Tensor<B, 5>) -> Tensor<B, 5> {
    type FT = FloatElem<TestBackend>;
    fn initializer_default() {
    fn initializer_zeros() {
    fn display() {
    fn input_channels_mismatch() {
/// Configuration to create a [deformable 2D convolution](DeformConv2d) layer, using the [init function](DeformConv2dConfig::init).
pub struct DeformConv2dConfig {
/// The number of channels.
/// The size of the kernel.
/// The stride of the convolution.
/// Spacing between kernel elements.
/// Controls the connections between input and output channels.
/// Offset groups.
/// The padding configuration.
///
/// ### Warning
/// Only symmetric padding is currently supported. As such, using `Same` padding with an even kernel
/// size is not supported as it will not produce the same output size.
/// If bias should be added to the output.
/// The type of function used to initialize neural network parameters
/// Applies a deformable 2D convolution over input tensors.
///
/// Should be created with [DeformConv2dConfig].
pub struct DeformConv2d<B: Backend> {
/// Tensor of shape `[channels_out, channels_in / groups, kernel_size_1, kernel_size_2]`
/// Tensor of shape `[channels_out]`
/// Stride of the convolution.
/// Size of the kernel.
/// Spacing between kernel elements.
/// Controls the connections between input and output channels.
/// Offset groups.
/// The padding configuration.
impl DeformConv2dConfig {
/// Initialize a new [DeformConv2d](DeformConv2d) module.
    pub fn init<B: Backend>(&self, device: &B::Device) -> DeformConv2d<B> {
impl<B: Backend> ModuleDisplay for DeformConv2d<B> {
    fn custom_settings(&self) -> Option<DisplaySettings> {
    fn custom_content(&self, content: Content) -> Option<Content> {
// Since padding does not implement ModuleDisplay, we need to format it manually.
// Format the stride, kernel_size and dilation as strings, formatted as arrays instead of indexed.
impl<B: Backend> DeformConv2d<B> {
/// Applies the forward pass on the input tensor.
///
/// See [deform_conv2d](burn::tensor::module::deform_conv2d) for more information.
///
/// # Shapes
///
/// - input: `[batch_size, channels_in, height_in, width_in]`
/// - offset: `[batch_size, 2 * offset_groups * kernel_height * kernel_width, height_out, width_out]`
/// - mask: `[batch_size, offset_groups * kernel_height * kernel_width, height_out, width_out]`
/// - output: `[batch_size, channels_out, height_out, width_out]`
    pub fn forward(
        &self,
        input: Tensor<B, 4>,
        offset: Tensor<B, 4>,
        mask: Option<Tensor<B, 4>>,
    ) -> Tensor<B, 4> {
    type FT = FloatElem<TestBackend>;
    fn initializer_default() {
    fn initializer_zeros() {
    fn initializer_fan_out() {
// test that fan_out is passed to `init_with()`
    fn initializer_fan_with_groups_is_valid() {
    fn channels_with_groups_is_invalid() {
    fn same_with_even_kernel_is_invalid() {
    fn display() {
    fn input_channels_mismatch() {
/// Configuration to create a [Dropout](Dropout) layer using the [init function](DropoutConfig::init).
pub struct DropoutConfig {
/// The probability of randomly zeroes some elements of the input tensor during training.
/// Set at random some elements of the input tensor to zero during training.
///
/// This is an effective regularization technique as describe in the paper
/// [Improving neural networks by preventing co-adaptation of feature detectors](https://arxiv.org/abs/1207.0580).
///
/// The input is also scaled during training to `1 / (1 - prob_keep)`.
///
/// Should be created with [DropoutConfig].
pub struct Dropout {
/// The probability of randomly zeroes some elements of the input tensor during training.
impl DropoutConfig {
/// Initialize a new [dropout](Dropout) module.
    pub fn init(&self) -> Dropout {
impl Dropout {
/// Applies the forward pass on the input tensor.
///
/// See [Dropout](Dropout) for more information.
///
/// # Shapes
///
/// - input: `[..., any]`
/// - output: `[..., any]`
    pub fn forward<B: Backend, const D: usize>(&self, input: Tensor<B, D>) -> Tensor<B, D> {
impl ModuleDisplay for Dropout {
    fn custom_settings(&self) -> Option<DisplaySettings> {
    fn custom_content(&self, content: Content) -> Option<Content> {
    fn with_ad_backend_should_mark_input() {
    fn without_ad_backend_should_not_change_input() {
    fn display() {
    fn dropout_prob_invalid() {
/// Configuration to create an [Embedding](Embedding) layer using the [init function](EmbeddingConfig::init).
pub struct EmbeddingConfig {
/// The number of embedding vectors.
/// The size of each vector.
/// The type of function used to initialize neural network parameters
/// Lookup table to store a fix number of vectors.
///
/// Should be created with [EmbeddingConfig].
pub struct Embedding<B: Backend> {
/// The learnable weights of the module of shape `[n_embedding, d_model]` initialized
/// from a normal distribution `N(0, 1)`.
impl<B: Backend> ModuleDisplay for Embedding<B> {
    fn custom_settings(&self) -> Option<DisplaySettings> {
    fn custom_content(&self, content: Content) -> Option<Content> {
impl EmbeddingConfig {
/// Initialize a new [embedding](Embedding) module.
    pub fn init<B: Backend>(&self, device: &B::Device) -> Embedding<B> {
impl<B: Backend> Embedding<B> {
/// Applies the forward pass on the input tensor.
///
/// See also [embedding](burn::tensor::module::embedding).
///
/// # Shapes
///
/// - input: `[batch_size, seq_length]`
/// - output: `[batch_size, seq_length, d_model]`
    pub fn forward(&self, input: Tensor<B, 2, Int>) -> Tensor<B, 3> {
    type FT = FloatElem<TestBackend>;
    fn initializer_zeros() {
    fn display() {
/// Configuration for the 1D interpolation module.
///
/// This struct defines the configuration options for the 1D interpolation operation.
/// It allows specifying the output size, scale factor, and interpolation mode.
pub struct Interpolate1dConfig {
/// Output size of the interpolated tensor.
/// If specified, this takes precedence over `scale_factor`.
/// Scale factor for resizing the input tensor.
/// This is used when `output_size` is not specified.
/// Interpolation mode to use for resizing.
/// Determines how the output values are calculated.
/// Interpolate module for resizing 1D tensors with shape [N, C, L].
///
/// This struct represents a 1D interpolation module that can resize tensors
/// using various interpolation methods. It provides flexibility in specifying
/// either an output size or a scale factor for resizing, along with options
/// for the interpolation mode.
///
/// The module can be used to upsample or downsample 1D tensors, preserving the
/// number of channels and batch size while adjusting the length dimension.
///
/// The module can be created using the [Interpolate1dConfig] struct and the
/// `init` method, which returns an instance of the [Interpolate1d] struct.
pub struct Interpolate1d {
/// Output size of the interpolated tensor
/// Scale factor for resizing the input tensor
/// Interpolation mode used for resizing
impl Interpolate1dConfig {
/// Initialize the interpolation module
    pub fn init(self) -> Interpolate1d {
impl Interpolate1d {
/// Performs the forward pass of the 1D interpolation module
///
/// # Arguments
///
/// * `input` - Input tensor with shape [N, C, L]
///
/// # Returns
///
/// Resized tensor with shape [N, C, L'], where L' is determined by
/// the output_size or scale_factor specified in the module configuration
///
/// # Example
///
/// ```ignore
/// let input = Tensor::<Backend, 3>::random([1, 3, 64], Distribution::Uniform(0.0, 1.0), &device);
/// let interpolate = Interpolate1dConfig::new()
///     .with_output_size(Some(128))
///     .init();
/// let output = interpolate.forward(input);
/// assert_eq!(output.dims(), [1, 3, 128]);
/// ```
    pub fn forward<B: Backend>(&self, input: Tensor<B, 3>) -> Tensor<B, 3> {
// Use the interpolate operation to resize the temporal input tensor
// by adding a new dimension for the interpolation axis
/// Calculate output size based on input dimensions, output size, and scale factor
///
/// # Arguments
///
/// * `input_dims` - Input dimensions of the tensor
/// * `output_size` - Output size for the interpolated tensor
/// * `scale_factor` - Scale factor for resizing the tensor
///
/// # Returns
///
/// Output size for the interpolated tensor
///
/// # Panics
///
/// Panics if neither output_size nor scale_factor is provided
/// or if the scale factor is too large
fn calculate_output_size(
    input_dims: [usize;
// Use provided
// Calculate output size based on scale factor
impl ModuleDisplay for Interpolate1d {
    fn custom_settings(&self) -> Option<DisplaySettings> {
    fn custom_content(&self, content: Content) -> Option<Content> {
    fn test_calculate_output_size() {
    fn test_panic() {
    fn test_large_scale_factor() {
    fn test_module() {
// Test with output_size
// Test with scale_factor
// Test with different interpolation mode
    fn display() {
/// Configuration for the 2D interpolation module.
///
/// This struct defines the configuration options for the 2D interpolation operation.
/// It allows specifying the output size, scale factor, and interpolation mode.
pub struct Interpolate2dConfig {
/// Output size of the interpolated tensor.
/// If specified, this takes precedence over `scale_factor`.
/// Scale factor for resizing the input tensor.
/// This is used when `output_size` is not specified.
/// Interpolation mode to use for resizing.
/// Determines how the output values are calculated.
/// Interpolate module for resizing tensors with shape [N, C, H, W].
///
/// This struct represents an interpolation module that can resize tensors
/// using various interpolation methods. It provides flexibility in specifying
/// either an output size or a scale factor for resizing, along with options
/// for the interpolation mode.
///
/// The module can be used to upsample or downsample tensors, preserving the
/// number of channels and batch size while adjusting the height and width
/// dimensions.
///
/// The module can be created using the [Interpolate2dConfig] struct and the
/// `init` method, which returns an instance of the [Interpolate2d] struct.
pub struct Interpolate2d {
/// Output size of the interpolated tensor
/// Scale factor for resizing the input tensor
/// Interpolation mode used for resizing
impl Interpolate2dConfig {
/// Initialize the interpolation module
    pub fn init(self) -> Interpolate2d {
impl Interpolate2d {
/// Performs the forward pass of the interpolation module
///
/// # Arguments
///
/// * `input` - Input tensor with shape [N, C, H, W]
///
/// # Returns
///
/// Resized tensor with shape [N, C, H', W'], where H' and W' are determined by
/// the output_size or scale_factor specified in the module configuration
///
/// # Example
///
/// ```ignore
/// let input = Tensor::<Backend, 2>::random([1, 3, 64, 64], Distribution::Uniform(0.0, 1.0), &device);
/// let interpolate = Interpolate2dConfig::new()
///     .with_output_size(Some([128, 128]))
///     .init();
/// let output = interpolate.forward(input);
/// assert_eq!(output.dims(), [1, 3, 128, 128]);
/// ```
    pub fn forward<B: Backend>(&self, input: Tensor<B, 4>) -> Tensor<B, 4> {
/// Calculates the output size for tensor interpolation.
///
/// # Arguments
///
/// * `input_dims` - The dimensions of the input tensor [N, C, H, W].
/// * `output_size` - Optional desired output size [H', W'].
/// * `scale_factor` - Optional scale factor for height and width [scale_h, scale_w].
///
/// # Returns
///
/// A tuple [H', W'] representing the calculated output size.
///
/// # Panics
///
/// Panics if neither `output_size` nor `scale_factor` is provided,
/// or if the scale factor results in dimensions exceeding usize::MAX.
fn calculate_output_size(
    input_dims: [usize;
// Use provided
// Calculate output size based on scale factor
impl ModuleDisplay for Interpolate2d {
    fn custom_settings(&self) -> Option<DisplaySettings> {
    fn custom_content(&self, content: Content) -> Option<Content> {
    fn test_calculate_output_size() {
    fn test_missing_params() {
    fn test_infinite_height() {
    fn test_infinite_width() {
    fn test_module() {
// Test with output_size
// Test with scale_factor
// Test with different interpolation mode
    fn display() {
/// Algorithm used for downsampling and upsampling
///
/// This enum defines different interpolation modes for resampling data.
pub enum InterpolateMode {
/// Nearest-neighbor interpolation
///
/// This mode selects the value of the nearest sample point for each output pixel.
/// It is applicable for both temporal and spatial data.
/// Linear interpolation
///
/// This mode calculates the output value using linear
/// interpolation between nearby sample points.
///
/// It is applicable for both temporal and spatial data.
/// Cubic interpolation
///
/// This mode uses cubic interpolation to calculate the output value
/// based on surrounding sample points.
///
/// It is applicable for both temporal and spatial data and generally
/// provides smoother results than linear interpolation.
impl From<InterpolateMode> for OpsInterpolateMode {
    fn from(mode: InterpolateMode) -> Self {
/// Configuration to create a [`Linear`] layer using the [init function](LinearConfig::init).
pub struct LinearConfig {
/// The size of the input features.
/// The size of the output features.
/// If a bias should be applied during the linear transformation.
/// The type of function used to initialize neural network parameters
/// The layout in which the linear parameters are stored.
/// The layout in which the linear parameters are stored.
///
/// This can have performance impacts.
pub enum LinearLayout {
/// Parameters are stored in Row major.
/// Parameters are stored in Col major.
/// Applies a linear transformation to the input tensor.
///
/// Should be created with [LinearConfig]
///
/// `O = IW + b`
pub struct Linear<B: Backend> {
/// Matrix of shape `[d_input, d_output]` initialized from a uniform distribution:
///     `U(-k, k)`, where `k = sqrt(1 / d_input)`
/// Vector of size `d_output` initialized from a uniform distribution:
///     `U(-k, k)`, where `k = sqrt(1 / d_input)`
impl LinearConfig {
/// Initialize a new [`Linear`] module.
    pub fn init<B: Backend>(&self, device: &B::Device) -> Linear<B> {
// The param is already transposed when init. We re-transpose to have
// [d_output, d_input] while saving.
// When loading from record we have to transpose.
// When loading from initialization, we have to transpose.
impl<B: Backend> Linear<B> {
/// Applies the forward pass on the input tensor.
///
/// # Arguments
///
/// - `input` - The input tensor of shape `[..., d_input]`.
///
/// # Shapes
///
/// - input: `[..., d_input]`
/// - output: `[..., d_output]`
///
/// # Returns
///
/// The transformed tensor of shape `[..., d_output]`.
    pub fn forward<const D: usize>(&self, input: Tensor<B, D>) -> Tensor<B, D> {
impl<B: Backend> ModuleDisplay for Linear<B> {
    fn custom_settings(&self) -> Option<DisplaySettings> {
    fn custom_content(&self, content: Content) -> Option<Content> {
    type FT = FloatElem<TestBackend>;
    fn initializer_default() {
    fn initializer_zeros() {
    fn test_linear_forward_no_bias() {
    fn test_linear_forward_with_bias() {
    fn test_linear_1d() {
    fn display() {
    fn layout() {
// We go through serialization to trigger the mappers..
// We go through serialization to trigger the mappers.
//
// The test will fail if the mapper is not correctly given to the module after loading a
// record.
    fn col_row_same_result() {
/// Attention module
/// Cache module
/// Convolution module
/// Pooling module
/// Transformer module
/// Interpolate module
/// [`BatchNorm`] Configuration.
///
/// Used to create a [`BatchNorm`] layer using the [`BatchNormConfig::init`].
pub struct BatchNormConfig {
/// The number of features.
/// A value required for numerical stability. Default: 1e-5
/// Momentum used to update the metrics. Default: 0.1
/// Applies Batch Normalization over a tensor.
///
/// Based upon the paper [Batch Normalization](https://arxiv.org/abs/1502.03167).
///
/// Assumes input tensor is of shape ``[batch_size, channels, ...]``.
///
/// `Y = norm(X) * γ + β`
///
/// Where:
/// - `X` is the input tensor
/// - `Y` is the output tensor
/// - `norm` is the normalization function
/// - `γ` is the learnable weight
/// - `β` is the learnable bias
///
/// Should be created using [`BatchNormConfig`].
pub struct BatchNorm<B: Backend> {
/// The learnable weight gamma.
/// The learnable weight beta.
/// The running mean.
/// The running variance.
/// Momentum used to update the metrics.
/// A value required for numerical stability.
impl BatchNormConfig {
/// Initializes a new [batch norm](BatchNorm) module.
    pub fn init<B: Backend>(&self, device: &B::Device) -> BatchNorm<B> {
impl<B: Backend> BatchNorm<B> {
/// Applies the forward pass on the input tensor.
///
/// See [`BatchNorm`] for more information.
///
/// # Shapes
///
/// - `input`: ``[batch_size, channels, ...]``
/// - `output`: ``[batch_size, channels, ...]``
///
/// # Panics
///
/// This function will panic if the input tensor has rank < 2.
    pub fn forward<const D: usize>(&self, input: Tensor<B, D>) -> Tensor<B, D> {
// Should be move to a compilation error when const generic support that kind of
// validation. https://github.com/rust-lang/rust/issues/76560
    fn forward_inference<const D: usize>(&self, input: Tensor<B, D>) -> Tensor<B, D> {
    fn forward_train<const D: usize>(&self, input: Tensor<B, D>) -> Tensor<B, D> {
    fn forward_shared<const D: usize>(
        &self,
        x: Tensor<B, D>,
        mean: Tensor<B, D>,
        var: Tensor<B, D>,
    ) -> Tensor<B, D> {
impl<B: Backend> ModuleDisplay for BatchNorm<B> {
    fn custom_settings(&self) -> Option<DisplaySettings> {
    fn custom_content(&self, content: Content) -> Option<Content> {
    type FT = FloatElem<TestAutodiffBackend>;
    fn batch_norm_forward_train() {
    fn batch_norm_forward_inference() {
    fn input_tensor<B: Backend>(device: &B::Device) -> Tensor<B, 3> {
    type FT = FloatElem<TestAutodiffBackend>;
    fn batch_norm_forward_train() {
    fn batch_norm_forward_inference() {
    fn batch_norm_running_mean() {
    fn batch_norm_running_var() {
    fn batch_norm_running_mean_inner_module() {
    fn batch_norm_grads() {
    fn input_tensor<B: Backend>(device: &B::Device) -> Tensor<B, 4> {
    fn display() {
/// Configuration to create a [GroupNorm](GroupNorm) layer using the [init function](GroupNormConfig::init).
pub struct GroupNormConfig {
/// The number of groups to separate the channels into
/// The number of channels expected in the input
/// A value required for numerical stability. Default: 1e-5
/// A boolean value that when set to `true`, this module has learnable
/// per-channel affine parameters initialized to ones (for weights)
/// and zeros (for biases). Default: `true`
/// Applies Group Normalization over a mini-batch of inputs as described in the paper [Group Normalization](https://arxiv.org/abs/1803.08494).
///
/// `Y = groupnorm(X) * γ + β`
///
/// Where:
/// - `X` is the input tensor
/// - `Y` is the output tensor
/// - `γ` is the learnable weight
/// - `β` is the learnable bias
///
/// Should be created using [GroupNormConfig](GroupNormConfig).
pub struct GroupNorm<B: Backend> {
/// The learnable weight
/// The learnable bias
/// The number of groups to separate the channels into
/// The number of channels expected in the input
/// A value required for numerical stability
/// A boolean value that when set to `true`, this module has learnable
impl<B: Backend> ModuleDisplay for GroupNorm<B> {
    fn custom_settings(&self) -> Option<DisplaySettings> {
    fn custom_content(&self, content: Content) -> Option<Content> {
impl GroupNormConfig {
/// Initialize a new [group norm](GroupNorm) module.
    pub fn init<B: Backend>(&self, device: &B::Device) -> GroupNorm<B> {
impl<B: Backend> GroupNorm<B> {
/// Applies the forward pass on the input tensor.
///
/// See [GroupNorm](GroupNorm) for more information.
///
/// # Shapes
///
/// - input: `[batch_size, num_channels, *]`
/// - output: `[batch_size, num_channels, *]`
    pub fn forward<const D: usize>(&self, input: Tensor<B, D>) -> Tensor<B, D> {
/// Applies Group Normalization over a mini-batch of inputs as described in the paper [Group Normalization](https://arxiv.org/abs/1803.08494).
///
/// `Y = groupnorm(X) * γ + β`
///
/// Where:
/// - `X` is the input tensor
/// - `Y` is the output tensor
/// - `γ` is the learnable weight
/// - `β` is the learnable bias
///
pub(crate) fn group_norm<B: Backend, const D: usize>(
    input: Tensor<B, D>,
    gamma: Option<Tensor<B, 1>>,
    beta: Option<Tensor<B, 1>>,
    num_groups: usize,
    epsilon: f64,
    affine: bool,
) -> Tensor<B, D> {
    type FT = FloatElem<TestBackend>;
    fn group_norm_forward_affine_false() {
    fn group_norm_forward_affine_true() {
    fn display() {
/// Configuration to create a [InstanceNorm](InstanceNorm) layer using the [init function](InstanceNormConfig::init).
pub struct InstanceNormConfig {
/// The number of channels expected in the input
/// A value required for numerical stability. Default: 1e-5
/// A boolean value that when set to `true`, this module has learnable
/// per-channel affine parameters initialized to ones (for weights)
/// and zeros (for biases). Default: `true`
/// Applies Instance Normalization over a tensor as described in the paper [Instance Normalization](https://arxiv.org/abs/1607.08022)
///
/// Should be created using [InstanceNormConfig](InstanceNormConfig).
pub struct InstanceNorm<B: Backend> {
/// The learnable weight
/// The learnable bias
/// The number of channels expected in the input
/// A value required for numerical stability
/// A boolean value that when set to `true`, this module has learnable
impl<B: Backend> ModuleDisplay for InstanceNorm<B> {
    fn custom_settings(&self) -> Option<DisplaySettings> {
    fn custom_content(&self, content: Content) -> Option<Content> {
impl InstanceNormConfig {
/// Initialize a new [instance norm](InstanceNorm) module.
    pub fn init<B: Backend>(&self, device: &B::Device) -> InstanceNorm<B> {
impl<B: Backend> InstanceNorm<B> {
/// Applies the forward pass on the input tensor.
///
/// See also [InstanceNormConfig](InstanceNormConfig) for more information.
///
/// # Shapes
///
/// - input: `[batch_size, num_channels, *]`
/// - output: `[batch_size, num_channels, *]`
    pub fn forward<const D: usize>(&self, input: Tensor<B, D>) -> Tensor<B, D> {
// Instance norm is equivalent to group norm when the number of groups is equal to the number of channels.
    type FT = FloatElem<TestBackend>;
    fn instance_norm_forward_affine_false() {
    fn instance_norm_forward_affine_true() {
    fn display() {
/// Configuration to create a [LayerNorm](LayerNorm) layer using the [init function](LayerNormConfig::init).
pub struct LayerNormConfig {
/// The size of the input features.
/// A value required for numerical stability. Default: 1e-5
/// Applies Layer Normalization over an input tensor as described in the paper [Layer Normalization](https://arxiv.org/abs/1607.06450).
///
/// `Y = norm(X) * γ + β`
///
/// Where:
/// - `X` is the input tensor
/// - `Y` is the output tensor
/// - `γ` is the learnable weight
/// - `β` is the learnable bias
///
/// Should be created using [LayerNormConfig](LayerNormConfig).
pub struct LayerNorm<B: Backend> {
/// The learnable weight.
/// The learnable bias.
/// A value required for numerical stability.
impl LayerNormConfig {
/// Initialize a new [layer norm](LayerNorm) module.
    pub fn init<B: Backend>(&self, device: &B::Device) -> LayerNorm<B> {
impl<B: Backend> LayerNorm<B> {
/// Applies the forward pass on the input tensor.
///
/// See the [LayerNorm](LayerNorm) documentation for more information.
///
/// # Shapes
///
/// - input: `[..., any, d_model]`
/// - output: `[..., any, d_model]`
    pub fn forward<const D: usize>(&self, input: Tensor<B, D>) -> Tensor<B, D> {
impl<B: Backend> ModuleDisplay for LayerNorm<B> {
    fn custom_settings(&self) -> Option<DisplaySettings> {
    fn custom_content(&self, content: Content) -> Option<Content> {
    type FT = FloatElem<TestBackend>;
    fn layer_norm_forward() {
    fn layer_norm_forward_large_epsilon() {
    fn layer_norm_backward() {
    fn display() {
//! # Normalization Layers
//!
//! Users who wish to provide an abstraction over swappable normalization
//! layers can use the [`Normalization`] wrapper, with support for:
//! * [`Normalization::Batch`] - [`BatchNorm`]
//! * [`Normalization::Group`] - [`GroupNorm`]
//! * [`Normalization::Instance`] - [`InstanceNorm`]
//! * [`Normalization::Layer`] - [`LayerNorm`]
//! * [`Normalization::Rms`] - [`RmsNorm`]
//!
//! [`NormalizationConfig`] can be used as a generic normalization policy:
//! * Construct a config with arbitrary input features (we suggest `0`).
//! * Clone and match that config to the target input layer,
//!   using the [`NormalizationConfig::with_num_features()`] method.
/// ['Normalization'] Configuration.
///
/// The enum is non-exhaustive to prepare for future additions.
///
/// Can be used as a generic configuration for normalization layers:
/// * Construct a config with arbitrary input features (we suggest `0`).
/// * Clone and match that config to the target input layer,
///   using the [`NormalizationConfig::with_num_features()`] method.
pub enum NormalizationConfig {
/// ['BatchNorm'] Configuration.
/// ['GroupNorm'] Configuration.
/// ['InstanceNorm'] Configuration.
/// ['LayerNorm'] Configuration.
/// ['RmsNorm'] Configuration.
impl From<BatchNormConfig> for NormalizationConfig {
    fn from(config: BatchNormConfig) -> Self {
impl From<GroupNormConfig> for NormalizationConfig {
    fn from(config: GroupNormConfig) -> Self {
impl From<InstanceNormConfig> for NormalizationConfig {
    fn from(config: InstanceNormConfig) -> Self {
impl From<LayerNormConfig> for NormalizationConfig {
    fn from(config: LayerNormConfig) -> Self {
impl From<RmsNormConfig> for NormalizationConfig {
    fn from(config: RmsNormConfig) -> Self {
impl NormalizationConfig {
/// Initialize a ['Norm'] layer.
    pub fn init<B: Backend>(&self, device: &B::Device) -> Normalization<B> {
/// Set the number of features.
    pub fn with_num_features(self, num_features: usize) -> Self {
/// Get the number of features.
    pub fn num_features(&self) -> usize {
/// Normalization Layer Wrapper
///
/// Provides support for built-in ``burn::nn::norm`` norm layers:
/// * [`Normalization::Batch`] - [`BatchNorm`]
/// * [`Normalization::Group`] - [`GroupNorm`]
/// * [`Normalization::Instance`] - [`InstanceNorm`]
/// * [`Normalization::Layer`] - [`LayerNorm`]
/// * [`Normalization::Rms`] - [`RmsNorm`]
///
/// The enum is non-exhaustive, to prepare for future additions.
pub enum Normalization<B: Backend> {
/// [`BatchNorm`] layer.
/// [`GroupNorm`] layer.
/// ['InstanceNorm'] layer.
/// [`LayerNorm`] layer.
/// ['RmsNorm'] layer.
impl<B: Backend> From<BatchNorm<B>> for Normalization<B> {
    fn from(layer: BatchNorm<B>) -> Self {
impl<B: Backend> From<GroupNorm<B>> for Normalization<B> {
    fn from(layer: GroupNorm<B>) -> Self {
impl<B: Backend> From<InstanceNorm<B>> for Normalization<B> {
    fn from(layer: InstanceNorm<B>) -> Self {
impl<B: Backend> From<LayerNorm<B>> for Normalization<B> {
    fn from(layer: LayerNorm<B>) -> Self {
impl<B: Backend> From<RmsNorm<B>> for Normalization<B> {
    fn from(layer: RmsNorm<B>) -> Self {
impl<B: Backend> Normalization<B> {
/// Applies normalization to a tensor.
///
/// The normalization contract depends upon the wrapped norm layer;
/// but all norm layers assume an input of at least rank 2;
/// and produce an output of the same rank and shape.
    pub fn forward<const D: usize>(&self, input: Tensor<B, D>) -> Tensor<B, D> {
/// Get the number of features.
    pub fn num_features(&self) -> usize {
    fn test_match_feature_size() {
    fn test_batch_norm() {
        type B = TestAutodiffBackend;
    fn test_group_norm() {
        type B = TestAutodiffBackend;
    fn test_instance_norm() {
        type B = TestAutodiffBackend;
    fn test_layer_norm() {
        type B = TestAutodiffBackend;
    fn test_rms_norm() {
        type B = TestAutodiffBackend;
/// Configuration to create a [RMS Norm](RmsNorm) layer using the [init function](RmsNormConfig::init).
pub struct RmsNormConfig {
/// The size of the input features.
/// A value required for numerical stability. Default: 1e-5
impl RmsNormConfig {
/// Initialize a new [RMS Norm](RmsNorm) module.
///
/// # Panics
///
/// Panics if `epsilon` is not positive.
    pub fn init<B: Backend>(&self, device: &B::Device) -> RmsNorm<B> {
/// Applies RMS Normalization over an input tensor along the last dimension.
///
/// `Y = X / sqrt(mean(X^2) + eps) * gamma`
///
/// Where:
/// - `X` is the input tensor
/// - `Y` is the output tensor
/// - `gamma` is the learnable weight
/// - `mean` is the mean operation
/// - `eps` is a small value to avoid division by zero.
///
/// Should be created using the [RmsNormConfig](RmsNormConfig) configuration.
pub struct RmsNorm<B: Backend> {
/// The learnable parameter to scale the normalized tensor
/// A value required for numerical stability
impl<B: Backend> RmsNorm<B> {
/// Applies the forward pass on the input tensor.
///
/// See the [RmsNorm](RmsNorm) documentation for more information.
///
/// # Shapes
///
/// - input: `[..., any, d_model]`
/// - output: `[..., any, d_model]`
    pub fn forward<const D: usize>(&self, x: Tensor<B, D>) -> Tensor<B, D> {
// Calculate the root-mean-square norm of the input tensor along the last dimension
impl<B: Backend> ModuleDisplay for RmsNorm<B> {
    fn custom_settings(&self) -> Option<DisplaySettings> {
    fn custom_content(&self, content: Content) -> Option<Content> {
    type FT = FloatElem<TestBackend>;
    fn rms_norm_forward() {
    fn display() {
/// Configuration to create a [1D adaptive avg pooling](AdaptiveAvgPool1d) layer using the [init function](AdaptiveAvgPool1dConfig::init).
pub struct AdaptiveAvgPool1dConfig {
/// The size of the output.
/// Applies a 1D adaptive avg pooling over input tensors.
///
/// Should be created with [AdaptiveAvgPool1dConfig].
pub struct AdaptiveAvgPool1d {
/// The size of the output.
impl ModuleDisplay for AdaptiveAvgPool1d {
    fn custom_settings(&self) -> Option<DisplaySettings> {
    fn custom_content(&self, content: Content) -> Option<Content> {
impl AdaptiveAvgPool1dConfig {
/// Initialize a new [adaptive avg pool 1d](AdaptiveAvgPool1d) module.
    pub fn init(&self) -> AdaptiveAvgPool1d {
impl AdaptiveAvgPool1d {
/// Applies the forward pass on the input tensor.
///
/// See [adaptive_avg_pool1d](burn::tensor::module::adaptive_avg_pool1d) for more information.
///
/// # Shapes
///
/// - input: `[batch_size, channels, length]`
/// - output: `[batch_size, channels, length_out]`
    pub fn forward<B: Backend>(&self, input: Tensor<B, 3>) -> Tensor<B, 3> {
    fn display() {
/// Configuration to create a [2D adaptive avg pooling](AdaptiveAvgPool2d) layer using the [init function](AdaptiveAvgPool2dConfig::init).
pub struct AdaptiveAvgPool2dConfig {
/// The size of the output.
/// Applies a 2D adaptive avg pooling over input tensors.
///
/// Should be created with [AdaptiveAvgPool2dConfig].
pub struct AdaptiveAvgPool2d {
/// The size of the output.
impl ModuleDisplay for AdaptiveAvgPool2d {
    fn custom_settings(&self) -> Option<DisplaySettings> {
    fn custom_content(&self, content: Content) -> Option<Content> {
impl AdaptiveAvgPool2dConfig {
/// Initialize a new [adaptive avg pool 2d](AdaptiveAvgPool2d) module.
    pub fn init(&self) -> AdaptiveAvgPool2d {
impl AdaptiveAvgPool2d {
/// Applies the forward pass on the input tensor.
///
/// See [adaptive_avg_pool2d](burn::tensor::module::adaptive_avg_pool2d) for more information.
///
/// # Shapes
///
/// - input: `[batch_size, channels, height_in, width_in]`
/// - output: `[batch_size, channels, height_out, width_out]`
    pub fn forward<B: Backend>(&self, input: Tensor<B, 4>) -> Tensor<B, 4> {
    fn display() {
/// Configuration to create a [1D avg pooling](AvgPool1d) layer using the [init function](AvgPool1dConfig::init).
pub struct AvgPool1dConfig {
/// The size of the kernel.
/// The stride.
/// The padding configuration.
///
/// ### Warning
/// Only symmetric padding is currently supported. As such, using `Same` padding with an even kernel
/// size is not supported as it will not produce the same output size.
/// If the padding is counted in the denominator when computing the average.
/// Applies a 1D avg pooling over input tensors.
///
/// Should be created with [AvgPool1dConfig](AvgPool1dConfig).
///
/// # Remarks
///
/// The zero-padding values will be included in the calculation
/// of the average. This means that the zeros are counted as
/// legitimate values, and they contribute to the denominator
/// when calculating the average. This is equivalent to
/// `torch.nn.AvgPool2d` with `count_include_pad=True`.
pub struct AvgPool1d {
/// The stride.
/// The size of the kernel.
/// The padding configuration.
/// If the padding is counted in the denominator when computing the average.
impl ModuleDisplay for AvgPool1d {
    fn custom_settings(&self) -> Option<DisplaySettings> {
    fn custom_content(&self, content: Content) -> Option<Content> {
impl AvgPool1dConfig {
/// Initialize a new [avg pool 1d](AvgPool1d) module.
    pub fn init(&self) -> AvgPool1d {
impl AvgPool1d {
/// Applies the forward pass on the input tensor.
///
/// See [avg_pool1d](burn::tensor::module::avg_pool1d) for more information.
///
/// # Shapes
///
/// - input: `[batch_size, channels, length_in]`
/// - output: `[batch_size, channels, length_out]`
    pub fn forward<B: Backend>(&self, input: Tensor<B, 3>) -> Tensor<B, 3> {
    fn same_with_even_kernel_is_invalid() {
    fn display() {
    fn default_strides_match_kernel_size(#[case] kernel_size: usize) {
/// Configuration to create a [2D avg pooling](AvgPool2d) layer using the [init function](AvgPool2dConfig::init).
pub struct AvgPool2dConfig {
/// The size of the kernel.
/// The strides.
/// The padding configuration.
///
/// ### Warning
/// Only symmetric padding is currently supported. As such, using `Same` padding with an even kernel
/// size is not supported as it will not produce the same output size.
/// If the padding is counted in the denominator when computing the average.
/// Applies a 2D avg pooling over input tensors.
///
/// Should be created with [AvgPool2dConfig](AvgPool2dConfig).
///
/// # Remarks
///
/// The zero-padding values will be included in the calculation
/// of the average. This means that the zeros are counted as
/// legitimate values, and they contribute to the denominator
/// when calculating the average. This is equivalent to
/// `torch.nn.AvgPool2d` with `count_include_pad=True`.
pub struct AvgPool2d {
/// Stride of the pooling.
/// Size of the kernel.
/// Padding configuration.
/// If the padding is counted in the denominator when computing the average.
impl ModuleDisplay for AvgPool2d {
    fn custom_settings(&self) -> Option<DisplaySettings> {
    fn custom_content(&self, content: Content) -> Option<Content> {
impl AvgPool2dConfig {
/// Initialize a new [avg pool 2d](AvgPool2d) module.
    pub fn init(&self) -> AvgPool2d {
impl AvgPool2d {
/// Applies the forward pass on the input tensor.
///
/// See [avg_pool2d](burn::tensor::module::avg_pool2d) for more information.
///
/// # Shapes
///
/// - input: `[batch_size, channels, height_in, width_in]`
/// - output: `[batch_size, channels, height_out, width_out]`
    pub fn forward<B: Backend>(&self, input: Tensor<B, 4>) -> Tensor<B, 4> {
    fn same_with_even_kernel_is_invalid() {
    fn display() {
    fn default_strides_match_kernel_size(#[case] kernel_size: [usize;
/// Configuration to create a [1D max pooling](MaxPool1d) layer using the [init function](MaxPool1dConfig::init).
pub struct MaxPool1dConfig {
/// The size of the kernel.
/// The stride.
/// The padding configuration.
///
/// ### Warning
/// Only symmetric padding is currently supported. As such, using `Same` padding with an even kernel
/// size is not supported as it will not produce the same output size.
/// The dilation.
/// Applies a 1D max pooling over input tensors.
///
/// Should be created with [MaxPool1dConfig](MaxPool1dConfig).
pub struct MaxPool1d {
/// The stride.
/// The size of the kernel.
/// The padding configuration.
/// The dilation.
impl ModuleDisplay for MaxPool1d {
    fn custom_settings(&self) -> Option<DisplaySettings> {
    fn custom_content(&self, content: Content) -> Option<Content> {
impl MaxPool1dConfig {
/// Initialize a new [max pool 1d](MaxPool1d) module.
    pub fn init(&self) -> MaxPool1d {
impl MaxPool1d {
/// Applies the forward pass on the input tensor.
///
/// See [max_pool1d](burn::tensor::module::max_pool1d) for more information.
///
/// # Shapes
///
/// - input: `[batch_size, channels, length_in]`
/// - output: `[batch_size, channels, length_out]`
    pub fn forward<B: Backend>(&self, input: Tensor<B, 3>) -> Tensor<B, 3> {
    fn same_with_even_kernel_is_invalid() {
    fn display() {
    fn default_strides_match_kernel_size(#[case] kernel_size: usize) {
/// Configuration to create a [2D max pooling](MaxPool2d) layer using the [init function](MaxPool2dConfig::init).
pub struct MaxPool2dConfig {
/// The size of the kernel.
/// The strides.
/// The padding configuration.
///
/// ### Warning
/// Only symmetric padding is currently supported. As such, using `Same` padding with an even kernel
/// size is not supported as it will not produce the same output size.
/// The dilation.
/// Applies a 2D max pooling over input tensors.
///
/// Should be created with [MaxPool2dConfig](MaxPool2dConfig).
pub struct MaxPool2d {
/// The strides.
/// The size of the kernel.
/// The padding configuration.
/// The dilation.
impl ModuleDisplay for MaxPool2d {
    fn custom_settings(&self) -> Option<DisplaySettings> {
    fn custom_content(&self, content: Content) -> Option<Content> {
impl MaxPool2dConfig {
/// Initialize a new [max pool 2d](MaxPool2d) module.
    pub fn init(&self) -> MaxPool2d {
impl MaxPool2d {
/// Applies the forward pass on the input tensor.
///
/// See [max_pool2d](burn::tensor::module::max_pool2d) for more information.
///
/// # Shapes
///
/// - input: `[batch_size, channels, height_in, width_in]`
/// - output: `[batch_size, channels, height_out, width_out]`
    pub fn forward<B: Backend>(&self, input: Tensor<B, 4>) -> Tensor<B, 4> {
    fn same_with_even_kernel_is_invalid() {
    fn display() {
    fn default_strides_match_kernel_size(#[case] kernel_size: [usize;
/// Configuration to create a [PositionalEncoding](PositionalEncoding) layer using the [init function](PositionalEncodingConfig::init).
pub struct PositionalEncodingConfig {
/// Maximum sequence size to use.
/// The size of each vector.
/// Max time scale to use.
/// Positional encoding layer for transformer models.
///
/// This layer adds positional information to the input embeddings, allowing the transformer model
/// to take into account the order of the sequence. The positional encoding is added to the input
/// embeddings by computing a set of sinusoidal functions with different frequencies and phases.
///
/// Sinusoids are used for positional embedding introduced in
/// [Attention is all you need](https://arxiv.org/abs/1706.03762).
///
/// The reference implementation can be found here:
/// [LANGUAGE MODELING WITH NN.TRANSFORMER AND TORCHTEXT
/// ](https://pytorch.org/tutorials/beginner/transformer_tutorial.html)
///
/// Should be created using [PositionalEncodingConfig]
pub struct PositionalEncoding<B: Backend> {
/// The sinusoids used to add positional information to the input embeddings.
/// The maximum sequence size to use.
/// Max time scale to use.
impl<B: Backend> ModuleDisplay for PositionalEncoding<B> {
    fn custom_settings(&self) -> Option<DisplaySettings> {
    fn custom_content(&self, content: Content) -> Option<Content> {
impl PositionalEncodingConfig {
/// Initialize a new [PositionalEncoding](PositionalEncoding) module.
    pub fn init<B: Backend>(&self, device: &B::Device) -> PositionalEncoding<B> {
impl<B: Backend> PositionalEncoding<B> {
/// Applies the forward pass on the input tensor by adding the sinusoids to the input.
///
/// # Shapes
///
/// * input: [batch_size, seq_length, d_model]
/// * output: [batch_size, seq_length, d_model]
///
///
/// # Panics
///
/// * Panics if the input sequence length is greater than the maximum sequence size.
/// * Panics if the input d_model is not equal to the d_model of the sinusoids.
    pub fn forward(&self, input: Tensor<B, 3>) -> Tensor<B, 3> {
/// Returns sinusoids for positional embedding introduced in
/// [Attention is all you need](https://arxiv.org/abs/1706.03762).
///
/// The reference implementation can be found here:
/// [LANGUAGE MODELING WITH NN.TRANSFORMER AND TORCHTEXT
/// ](https://pytorch.org/tutorials/beginner/transformer_tutorial.html)
///
/// # Arguments
///
/// * `length` - The length of the sequence.
/// * `d_model` - The size of each vector.
/// * `max_timescale` - The maximum time scale to use.
///
/// # Returns
///
/// A tensor of shape [length, d_model] containing the sinusoids.
pub fn generate_sinusoids<B: Backend>(
    length: usize,
    d_model: usize,
    max_timescale: usize,
    device: &B::Device,
) -> Tensor<B, 2> {
// Calculate the increment for the logarithmic timescale
// Create a vector to hold the sinusoids
// Loop over each position in the sequence
// Create a vector to hold the sinusoids for this position
// Loop over each dimension of the sinusoids
// Calculate the division term for this dimension
// Calculate the sine and cosine values for this dimension and position
// Add the sinusoids for this position to the vector
// Convert the sinusoids to a tensor and return it
    type FT = FloatElem<TestBackend>;
    fn test_module() {
// expected to broadcast
// Use a tensor of zeros as input for easy verification of the output
// The output should be the sinusoids broadcasted to the input shape
    fn test_generate_sinusoids() {
// The values are taken from the pytorch reference implementation
    fn d_model_input_should_match() {
    fn input_length_should_be_less_than_max_len() {
    fn display() {
/// A GateController represents a gate in an LSTM cell. An
/// LSTM cell generally contains three gates: an input gate,
/// forget gate, and output gate. Additionally, cell gate
/// is just used to compute the cell state.
///
/// An Lstm gate is modeled as two linear transformations.
/// The results of these transformations are used to calculate
/// the gate's output.
pub struct GateController<B: Backend> {
/// Represents the affine transformation applied to input vector
/// Represents the affine transformation applied to the hidden state
impl<B: Backend> GateController<B> {
/// Initialize a new [gate_controller](GateController) module.
    pub fn new(
        d_input: usize,
        d_output: usize,
        bias: bool,
        initializer: Initializer,
        device: &B::Device,
    ) -> Self {
/// Helper function for performing weighted matrix product for a gate and adds
/// bias, if any.
///
///  Mathematically, performs `Wx*X + Wh*H + b`, where:
///     Wx = weight matrix for the connection to input vector X
///     Wh = weight matrix for the connection to hidden state H
///     X = input vector
///     H = hidden state
///     b = bias terms
    pub fn gate_product(&self, input: Tensor<B, 2>, hidden: Tensor<B, 2>) -> Tensor<B, 2> {
/// Used to initialize a gate controller with known weight layers,
/// allowing for predictable behavior. Used only for testing in
/// lstm.
    pub fn create_with_weights(
        d_input: usize,
        d_output: usize,
        bias: bool,
        initializer: Initializer,
        input_record: crate::LinearRecord<B>,
        hidden_record: crate::LinearRecord<B>,
    ) -> Self {
/// Configuration to create a [gru](Gru) module using the [init function](GruConfig::init).
pub struct GruConfig {
/// The size of the input features.
/// The size of the hidden state.
/// If a bias should be applied during the Gru transformation.
/// If reset gate should be applied after weight multiplication.
///
/// This configuration option controls how the reset gate is applied to the hidden state.
/// * `true` - (Default) Match the initial arXiv version of the paper [Learning Phrase Representations using RNN Encoder-Decoder for
///   Statistical Machine Translation (v1)](https://arxiv.org/abs/1406.1078v1) and apply the reset gate after multiplication by
///   the weights. This matches the behavior of [PyTorch GRU](https://pytorch.org/docs/stable/generated/torch.nn.GRU.html#torch.nn.GRU).
/// * `false` - Match the most recent revision of [Learning Phrase Representations using RNN Encoder-Decoder for Statistical Machine
///   Translation (v3)](https://arxiv.org/abs/1406.1078) and apply the reset gate before the weight multiplication.
///
/// The differing implementations can give slightly different numerical results and have different efficiencies. For more
/// motivation for why the `true` can be more efficient see [Optimizing RNNs with Differentiable Graphs](https://svail.github.io/diff_graphs).
///
/// To set this field to `false` use [`with_reset_after`](`GruConfig::with_reset_after`).
/// Gru initializer
/// The Gru (Gated recurrent unit) module. This implementation is for a unidirectional, stateless, Gru.
///
/// Introduced in the paper: [Learning Phrase Representations using RNN Encoder-Decoder for Statistical Machine Translation](https://arxiv.org/abs/1406.1078).
///
/// Should be created with [GruConfig].
pub struct Gru<B: Backend> {
/// The update gate controller.
/// The reset gate controller.
/// The new gate controller.
/// The size of the hidden state.
/// If reset gate should be applied after weight multiplication.
impl<B: Backend> ModuleDisplay for Gru<B> {
    fn custom_settings(&self) -> Option<DisplaySettings> {
    fn custom_content(&self, content: Content) -> Option<Content> {
impl GruConfig {
/// Initialize a new [gru](Gru) module.
    pub fn init<B: Backend>(&self, device: &B::Device) -> Gru<B> {
impl<B: Backend> Gru<B> {
/// Applies the forward pass on the input tensor. This GRU implementation
/// returns a state tensor with dimensions `[batch_size, sequence_length, hidden_size]`.
///
/// # Parameters
/// - batched_input: `[batch_size, sequence_length, input_size]`.
/// - state: An optional tensor representing an initial cell state with dimensions
///   `[batch_size, hidden_size]`. If none is provided, an empty state will be used.
///
/// # Returns
/// - output: `[batch_size, sequence_length, hidden_size]`
    pub fn forward(
        &self,
        batched_input: Tensor<B, 3>,
        state: Option<Tensor<B, 2>>,
    ) -> Tensor<B, 3> {
// u(pdate)g(ate) tensors
// Colloquially referred to as z(t)
// r(eset)g(ate) tensors
// Colloquially referred to as r(t)
// n(ew)g(ate) tensor
// Passed as input to new_gate
// Colloquially referred to as g(t)
// calculate linear interpolation between previous hidden state and candidate state:
// g(t) * (1 - z(t)) + z(t) * hidden_t
// (1 - z(t)) = -(z(t) - 1)
/// Helper function for performing weighted matrix product for a gate and adds
/// bias, if any, and optionally applies reset to hidden state.
///
///  Mathematically, performs `Wx*X + r .* (Wh*H + b)`, where:
///     Wx = weight matrix for the connection to input vector X
///     Wh = weight matrix for the connection to hidden state H
///     X = input vector
///     H = hidden state
///     b = bias terms
///     r = reset state
    fn gate_product(
        &self,
        input: &Tensor<B, 2>,
        hidden: &Tensor<B, 2>,
        reset: Option<&Tensor<B, 2>>,
        gate: &GateController<B>,
    ) -> Tensor<B, 2> {
    type FT = FloatElem<TestBackend>;
    fn init_gru<B: Backend>(reset_after: bool, device: &B::Device) -> Gru<B> {
        fn create_gate_controller<B: Backend>(
            weights: f32,
            biases: f32,
            d_input: usize,
            d_output: usize,
            bias: bool,
            initializer: Initializer,
            device: &B::Device,
        ) -> GateController<B> {
/// Test forward pass with simple input vector.
///
/// z_t = sigmoid(0.5*0.1 + 0.5*0) = 0.5125
/// r_t = sigmoid(0.6*0.1 + 0.*0) = 0.5150
/// g_t = tanh(0.7*0.1 + 0.7*0) = 0.0699
///
/// h_t = z_t * h' + (1 - z_t) * g_t = 0.0341
    fn tests_forward_single_input_single_feature() {
// Reset gate applied to hidden state before the matrix multiplication
// Reset gate applied to hidden state after the matrix multiplication
// override forward behavior
    fn tests_forward_seq_len_3() {
// Reset gate applied to hidden state before the matrix multiplication
// override forward behavior
    fn test_batched_forward_pass() {
    fn display() {
/// A LstmState is used to store cell state and hidden state in LSTM.
pub struct LstmState<B: Backend, const D: usize> {
/// The cell state.
/// The hidden state.
impl<B: Backend, const D: usize> LstmState<B, D> {
/// Initialize a new [LSTM State](LstmState).
    pub fn new(cell: Tensor<B, D>, hidden: Tensor<B, D>) -> Self {
/// Configuration to create a [Lstm](Lstm) module using the [init function](LstmConfig::init).
pub struct LstmConfig {
/// The size of the input features.
/// The size of the hidden state.
/// If a bias should be applied during the Lstm transformation.
/// Lstm initializer
/// The Lstm module. This implementation is for a unidirectional, stateless, Lstm.
///
/// Introduced in the paper: [Long Short-Term Memory](https://www.researchgate.net/publication/13853244).
///
/// Should be created with [LstmConfig].
pub struct Lstm<B: Backend> {
/// The input gate regulates which information to update and store in the cell state at each time step.
/// The forget gate is used to control which information to discard or keep in the memory cell at each time step.
/// The output gate determines which information from the cell state to output at each time step.
/// The cell gate is used to compute the cell state that stores and carries information through time.
/// The hidden state of the LSTM.
impl<B: Backend> ModuleDisplay for Lstm<B> {
    fn custom_settings(&self) -> Option<DisplaySettings> {
    fn custom_content(&self, content: Content) -> Option<Content> {
impl LstmConfig {
/// Initialize a new [lstm](Lstm) module.
    pub fn init<B: Backend>(&self, device: &B::Device) -> Lstm<B> {
impl<B: Backend> Lstm<B> {
/// Applies the forward pass on the input tensor. This LSTM implementation
/// returns the state for each element in a sequence (i.e., across seq_length) and a final state.
///
/// ## Parameters:
/// - batched_input: The input tensor of shape `[batch_size, sequence_length, input_size]`.
/// - state: An optional `LstmState` representing the initial cell state and hidden state.
///   Each state tensor has shape `[batch_size, hidden_size]`.
///   If no initial state is provided, these tensors are initialized to zeros.
///
/// ## Returns:
/// - output: A tensor represents the output features of LSTM. Shape: `[batch_size, sequence_length, hidden_size]`
/// - state: A `LstmState` represents the final states. Both `state.cell` and `state.hidden` have the shape
///   `[batch_size, hidden_size]`.
    pub fn forward(
        &self,
        batched_input: Tensor<B, 3>,
        state: Option<LstmState<B, 2>>,
    ) -> (Tensor<B, 3>, LstmState<B, 2>) {
    fn forward_iter<I: Iterator<Item = (Tensor<B, 3>, usize)>>(
        &self,
        input_timestep_iter: I,
        state: Option<LstmState<B, 2>>,
        batch_size: usize,
        seq_length: usize,
        device: &B::Device,
    ) -> (Tensor<B, 3>, LstmState<B, 2>) {
// f(orget)g(ate) tensors
// to multiply with cell state
// i(nput)g(ate) tensors
// o(output)g(ate) tensors
// c(ell)g(ate) tensors
// store the hidden state for this timestep
/// Configuration to create a [BiLstm](BiLstm) module using the [init function](BiLstmConfig::init).
pub struct BiLstmConfig {
/// The size of the input features.
/// The size of the hidden state.
/// If a bias should be applied during the BiLstm transformation.
/// BiLstm initializer
/// The BiLstm module. This implementation is for Bidirectional LSTM.
///
/// Introduced in the paper: [Framewise phoneme classification with bidirectional LSTM and other neural network architectures](https://www.cs.toronto.edu/~graves/ijcnn_2005.pdf).
///
/// Should be created with [BiLstmConfig].
pub struct BiLstm<B: Backend> {
/// LSTM for the forward direction.
/// LSTM for the reverse direction.
/// The size of the hidden state.
impl<B: Backend> ModuleDisplay for BiLstm<B> {
    fn custom_settings(&self) -> Option<DisplaySettings> {
    fn custom_content(&self, content: Content) -> Option<Content> {
impl BiLstmConfig {
/// Initialize a new [Bidirectional LSTM](BiLstm) module.
    pub fn init<B: Backend>(&self, device: &B::Device) -> BiLstm<B> {
impl<B: Backend> BiLstm<B> {
/// Applies the forward pass on the input tensor. This Bidirectional LSTM implementation
/// returns the state for each element in a sequence (i.e., across seq_length) and a final state.
///
/// ## Parameters:
/// - batched_input: The input tensor of shape `[batch_size, sequence_length, input_size]`.
/// - state: An optional `LstmState` representing the initial cell state and hidden state.
///   Each state tensor has shape `[2, batch_size, hidden_size]`.
///   If no initial state is provided, these tensors are initialized to zeros.
///
/// ## Returns:
/// - output: A tensor represents the output features of LSTM. Shape: `[batch_size, sequence_length, hidden_size * 2]`
/// - state: A `LstmState` represents the final forward and reverse states. Both `state.cell` and
///   `state.hidden` have the shape `[2, batch_size, hidden_size]`.
    pub fn forward(
        &self,
        batched_input: Tensor<B, 3>,
        state: Option<LstmState<B, 3>>,
    ) -> (Tensor<B, 3>, LstmState<B, 3>) {
// forward direction
// reverse direction
    type FT = FloatElem<TestBackend>;
    fn test_with_uniform_initializer() {
/// Test forward pass with simple input vector.
///
/// f_t = sigmoid(0.7*0.1 + 0.7*0) = sigmoid(0.07) = 0.5173928
/// i_t = sigmoid(0.5*0.1 + 0.5*0) = sigmoid(0.05) = 0.5123725
/// o_t = sigmoid(1.1*0.1 + 1.1*0) = sigmoid(0.11) = 0.5274723
/// c_t = tanh(0.9*0.1 + 0.9*0) = tanh(0.09) = 0.0892937
/// C_t = f_t * 0 + i_t * c_t = 0 + 0.5123725 * 0.0892937 = 0.04575243
/// h_t = o_t * tanh(C_t) = 0.5274723 * tanh(0.04575243) = 0.5274723 * 0.04568173 = 0.024083648
    fn test_forward_single_input_single_feature() {
        fn create_gate_controller(
            weights: f32,
            biases: f32,
            d_input: usize,
            d_output: usize,
            bias: bool,
            initializer: Initializer,
            device: &Device<TestBackend>,
        ) -> GateController<TestBackend> {
// single timestep with single feature
    fn test_batched_forward_pass() {
    fn test_batched_forward_pass_batch_of_one() {
    fn test_batched_backward_pass() {
// Asserts that the gradients exist and are non-zero
    fn test_bidirectional() {
        fn create_gate_controller<const D1: usize, const D2: usize>(
            input_weights: [[f32;
    fn display_lstm() {
    fn display_bilstm() {
/// Gated Recurrent Unit module.
/// Long Short-Term Memory module.
/// Configuration to create a [RotaryEncoding](RotaryEncoding) layer using the [init function](RotaryEncodingConfig::init).
pub struct RotaryEncodingConfig {
/// Maximum sequence length of input
/// Size of the input embedding or hidden dimension
/// Scaling factor for frequency computation. Defaults to 10000.0
impl RotaryEncodingConfig {
/// Initialize a new [RotaryEncoding](RotaryEncoding) module.
///
/// # Panics
///
/// Panics if the size of input embedding dimension is not even.
/// Panics if the theta parameter is not positive.
    pub fn init<B: Backend>(&self, device: &B::Device) -> RotaryEncoding<B> {
/// Initialize a new [RotaryEncoding](RotaryEncoding) module with a custom frequency scaling function.
/// This is useful to apply different RoPE extensions.
///
/// # Panics
///
/// Panics if the size of input embedding dimension is not even.
/// Panics if the theta parameter is not positive.
    pub fn init_with_frequency_scaling<B: Backend>(
        &self,
        scaling: impl Fn(Tensor<B, 1>) -> Tensor<B, 1>,
        device: &B::Device,
    ) -> RotaryEncoding<B> {
/// Initialize a new [RotaryEncoding](RotaryEncoding) module.
///
/// # Panics
///
/// Panics if the size of input embedding dimension is not even.
/// Panics if the theta parameter is not positive.
    fn initialize<B: Backend>(
        &self,
        scaling: impl Fn(Tensor<B, 1>) -> Tensor<B, 1>,
        device: &B::Device,
    ) -> RotaryEncoding<B> {
// Calculate the rotation frequencies for positional embeddings based on the formula
// `theta = 1 / (theta ^ (2i / d_model)) for i in [0..d_model/2]`
// Calculate (10000 ^ (2i / d_model)) by using the log base property `exp(log(10000) * (2i / d_model))`
// This is done since burn doesn't support exponentiation of scalar to tensor
/// A module that applies rotary positional encoding to a tensor.
/// Rotary Position Encoding or Embedding (RoPE), is a type of position embedding which encodes
/// absolute positional information with rotation matrix and naturally incorporates
/// explicit relative position dependency in self-attention formulation.
///
/// Introduced in the paper: [RoFormer: Enhanced Transformer with Rotary Position Embedding](https://arxiv.org/abs/2104.09864)
///
/// Should be created using [RotaryEncodingConfig].
pub struct RotaryEncoding<B: Backend> {
/// Complex frequency tensor of shape (max_sequence_length, d_model, 2) with real and imaginary components
// Essentially a cache of pre-computed RoPE values.
/// Frequency vector used to compute/apply the complex rotations.
impl<B: Backend> ModuleDisplay for RotaryEncoding<B> {
    fn custom_settings(&self) -> Option<DisplaySettings> {
    fn custom_content(&self, content: Content) -> Option<Content> {
impl<B: Backend> RotaryEncoding<B> {
/// Applies rotary positional encoding to a tensor of dimensions (..., seq_len, d_model)
///
/// # Arguments:
/// * `x` - Input tensor of shape (..., seq_len, d_model). Accommodate both 3D and 4D tensors
///   for (batch size, seq_len, hidden_dim) or (batch size, num_heads, seq_len, hidden_dim)
///   respectively.
///
/// # Returns:
/// Output tensor with the same shape as input tensor after applying rotary encoding.
///
/// # Panics
/// If the input tensor does not have at least 2 dimensions for sequence length and hidden dimension.
    pub fn forward<const D: usize>(&self, x: Tensor<B, D>) -> Tensor<B, D> {
/// Applies rotary positional encoding to a tensor of dimensions (..., seq_len, d_model)
///
/// # Arguments:
/// * `x` - Input tensor of shape (..., seq_len, d_model). Accommodate both 3D and 4D tensors
///   for (batch size, seq_len, hidden_dim) or (batch size, num_heads, seq_len, hidden_dim)
///   respectively.
/// * `start` - Sequence start position index.
///
/// # Returns:
/// Output tensor with the same shape as input tensor after applying rotary encoding.
///
/// # Panics
/// If the input tensor does not have at least 2 dimensions for sequence length and hidden dimension.
    pub fn apply<const D: usize>(&self, x: Tensor<B, D>, start: usize) -> Tensor<B, D> {
// Extract the sequence length and embedding dimension, other dimensions are kept generic
// to allow both 3D and 4D tensors i.e. batch_size or (batch_size, num_heads)
// Create a dummy tensor with signed ones based on the 2D rotation matrix
// [[cos, -sin], [sin, cos]]
// Rotate input using the frequency tensor. Slice the frequencies till input sequence length
// Sum the real and imaginary components to get output tensor and reshape to original shape
/// Shifts the pre-computed rotary frequency to cover a new range of positions.
///
/// This method updates the internal frequency tensor `freq_complex` to store
/// the rotary positional encodings for a new window of positions starting at `start`.
    pub fn shift(&mut self, start: usize) {
// Overwrite the whole buffer
// Shift the tail
// Compute the rest and assign
/// Computes the positional rotation frequencies (cosine and sine values) used in RoPE.
///
/// # Arguments
/// - `range`: Range of position indices `[start, end)`.
/// - `theta`: 1D tensor of shape `(d_model / 2)` containing base angular frequencies.
///
/// # Returns
/// Tensor of shape `(range.len(), d_model, 2)` containing `[cos, sin]` pairs for each position and frequency.
    fn compute_rotary_frequencies(range: Range<usize>, theta: Tensor<B, 1>) -> Tensor<B, 3> {
// Generate frequency values for positional embeddings
// Convert frequency values to complex numbers (polar form)
    type FT = FloatElem<TestBackend>;
    fn test_rotary_encoding_forward() {
// Input = [Batch size, Num of heads, Seq_len, d_model]
    fn test_rotary_encoding_3d() {
// Input = [Batch size, Num of heads, Seq_len, d_model]
// let input = input.unsqueeze::<4>();
    fn test_zero_input_rotary_encoding_forward() {
// Use a tensor of exact zeros as input. The output rotary embedding should be zeros as well
    fn test_valid_input_hidden_dim() {
// Hidden dimension must be even to be able to split into real and imaginary components
// for rotation
    fn test_rotary_encoding_frequencies() {
    fn apply_freq_scaling_by_parts<B: Backend>(freqs: Tensor<B, 1>) -> Tensor<B, 1> {
// Adapted from: https://github.com/meta-llama/llama-models/blob/main/models/llama3/reference_impl/model.py#L45
// if wavelen >= high_freq_wavelen
// (1 - smooth) * freq / scale_factor + smooth * freq
// if wavelen > low_freq_wavelen
// if wavelen < high_freq_wavelen
    fn test_rotary_encoding_with_frequency_scaling() {
    fn test_rotary_encoding_shift_full() {
// Input = [Batch size, Num of heads, Seq_len, d_model]
// Initializing for a bigger cache (e.g., max_seq_len = 10) should give the same result
// as using a smaller cache of pre-computed RoPE frequencies that are shifted to the same
// initial position
// start > 4 will perform a full re-compute
    fn test_rotary_encoding_shift() {
// Input = [Batch size, Num of heads, Seq_len, d_model]
// Initializing for a bigger cache (e.g., max_seq_len = 10) should give the same result
// as using a smaller cache of pre-computed RoPE frequencies that are shifted to the same
// initial position
// start < 4 will shift the (current_end - start) freqs and compute the rest
    fn test_rotary_encoding_shift_multiple() {
    fn test_rotary_encoding_shift_should_increase() {
// should be monotonically increasing
    fn display() {
/// Configuration to create a [Transformer Decoder](TransformerDecoder) layer using the [init function](TransformerDecoderConfig::init).
pub struct TransformerDecoderConfig {
/// The size of the model.
/// The size of the position-wise feed-forward network.
/// The number of attention heads.
/// The number of layers.
/// The dropout rate. Default: 0.1
/// Layer norm will be applied first instead of after the other modules.
/// Use "quiet softmax" instead of regular softmax.
///
/// - Usage may improve performance by allowing attention heads to deposit no information (if the sequence contains no information relevant to that head).
/// - Usage may reduce the entropy of weights in the model, enhancing quantization and compression.
///
/// Reference: <https://www.evanmiller.org/attention-is-off-by-one.html>
/// The type of function used to initialize neural network parameters
/// The transformer decoder module as describe in the paper [Attention Is All You Need](https://arxiv.org/abs/1706.03762).
///
/// # Params
///
/// - layers: transformer decoder layers with `d_model` input and output features.
///
/// Should be created using [TransformerDecoderConfig]
pub struct TransformerDecoder<B: Backend> {
/// Transformer decoder layers.
/// The size of the model.
/// The size of the position-wise feed-forward network.
/// The number of attention heads.
/// The number of layers.
/// The dropout rate. Default: 0.1
/// Layer norm will be applied first instead of after the other modules.
/// Use "quiet softmax" instead of regular softmax.
impl<B: Backend> ModuleDisplay for TransformerDecoder<B> {
    fn custom_settings(&self) -> Option<DisplaySettings> {
    fn custom_content(&self, content: Content) -> Option<Content> {
impl TransformerDecoderConfig {
/// Initialize a new [Transformer Decoder](TransformerDecoder) module.
    pub fn init<B: Backend>(&self, device: &B::Device) -> TransformerDecoder<B> {
/// [Transformer Decoder](TransformerDecoder) forward pass input argument.
pub struct TransformerDecoderInput<B: Backend> {
impl<B: Backend> TransformerDecoderInput<B> {
/// Create a [transformer decoder](TransformerDecoder) input argument.
    pub fn new(target: Tensor<B, 3>, memory: Tensor<B, 3>) -> Self {
/// Register the memory padding mask.
    pub fn memory_mask_pad(mut self, mask_pad: Tensor<B, 2, Bool>) -> Self {
/// Register the memory attention mask.
    pub fn memory_mask_attn(mut self, mask_attn: Tensor<B, 3, Bool>) -> Self {
/// Register the target padding mask.
    pub fn target_mask_pad(mut self, mask_pad: Tensor<B, 2, Bool>) -> Self {
/// Register the target attention mask.
    pub fn target_mask_attn(mut self, mask_attn: Tensor<B, 3, Bool>) -> Self {
/// [Transformer Decoder](TransformerDecoder) layer module.
pub struct TransformerDecoderLayer<B: Backend> {
struct TransformerDecoderLayerAutoregressiveCache<B: Backend> {
impl<B: Backend> TransformerDecoderLayerAutoregressiveCache<B> {
    fn empty() -> Self {
/// Autoregressive cache for the [Transformer Decoder](TransformerDecoder) layer.
///
/// To be used during inference when decoding tokens.
pub struct TransformerDecoderAutoregressiveCache<B: Backend> {
impl<B: Backend> TransformerDecoderAutoregressiveCache<B> {
    fn empty(num_layers: usize) -> Self {
impl<B: Backend> TransformerDecoderLayer<B> {
    fn new(config: &TransformerDecoderConfig, device: &B::Device) -> Self {
/// Applies the TransformerDecoder forward pass to the input tensor.
    fn forward(&self, mut input: TransformerDecoderInput<B>) -> TransformerDecoderInput<B> {
// Self attention residual path.
// Normalize.
// Self attention.
// Cross attention residual path.
// Normalize.
// Cross attention.
// Feed forward residual path.
// Normalize.
// Main path.
// Normalize.
    fn forward_autoregressive_inference(
        &self,
        mut input: TransformerDecoderInput<B>,
        cache: &mut TransformerDecoderLayerAutoregressiveCache<B>,
    ) -> TransformerDecoderInput<B> {
// Self attention residual path.
// Normalize.
// Self attention.
// Cross attention residual path.
// Normalize.
// Cross attention.
// Feed forward residual path.
// Normalize.
// Main path.
// Normalize.
impl<B: Backend> TransformerDecoder<B> {
/// Applies the forward pass.
    pub fn forward(&self, mut input: TransformerDecoderInput<B>) -> Tensor<B, 3> {
/// Applies the forward pass on the input using autoregressive cache.
    pub fn forward_autoregressive_inference(
        &self,
        mut input: TransformerDecoderInput<B>,
        cache: &mut TransformerDecoderAutoregressiveCache<B>,
    ) -> Tensor<B, 3> {
/// Create an empty autoregressive cache.
    pub fn new_autoregressive_cache(&self) -> TransformerDecoderAutoregressiveCache<B> {
    type FT = FloatElem<TestBackend>;
    fn test_autoregressive_norm_last() {
    fn test_autoregressive_norm_first() {
    fn test_autoregressive(config: TransformerDecoderConfig) {
// Normal forward using masking.
// Forward using the autoregressive cache.
// Greedy sampling
// Should produce the same tokens.
    fn display() {
/// Configuration to create a [Transformer Encoder](TransformerEncoder) layer using the [init function](TransformerEncoderConfig::init).
pub struct TransformerEncoderConfig {
/// The size of the model.
/// The size of the position-wise feed-forward network.
/// The number of attention heads.
/// The number of layers.
/// The dropout rate. Default: 0.1
/// Layer norm will be applied first instead of after the other modules.
/// Use "quiet softmax" instead of regular softmax.
///
/// - Usage may improve performance by allowing attention heads to deposit no information (if the sequence contains no information relevant to that head).
/// - Usage may reduce the entropy of weights in the model, enhancing quantization and compression.
///
/// Reference: <https://www.evanmiller.org/attention-is-off-by-one.html>
/// The type of function used to initialize neural network parameters
/// The transformer encoder module as describe in the paper [Attention Is All You Need](https://arxiv.org/abs/1706.03762).
///
/// # Params
///
/// - layers: transformer encoder layers with `d_model` input and output features.
///
/// Should be created using [TransformerEncoderConfig]
pub struct TransformerEncoder<B: Backend> {
/// The transformer encoder layers.
/// The size of the model.
/// The size of the position-wise feed-forward network.
/// The number of attention heads.
/// The number of layers.
/// The dropout rate. Default: 0.1
/// Layer norm will be applied first instead of after the other modules.
/// Use "quiet softmax" instead of regular softmax.
impl<B: Backend> ModuleDisplay for TransformerEncoder<B> {
    fn custom_settings(&self) -> Option<DisplaySettings> {
    fn custom_content(&self, content: Content) -> Option<Content> {
/// [Transformer Encoder](TransformerEncoder) forward pass input argument.
pub struct TransformerEncoderInput<B: Backend> {
impl<B: Backend> TransformerEncoderInput<B> {
/// Create a [transformer encoder](TransformerEncoder) input argument.
    pub fn new(tensor: Tensor<B, 3>) -> Self {
/// Register the padding mask.
    pub fn mask_pad(mut self, mask_pad: Tensor<B, 2, Bool>) -> Self {
/// Register the attention mask.
    pub fn mask_attn(mut self, mask_attn: Tensor<B, 3, Bool>) -> Self {
impl TransformerEncoderConfig {
/// Initialize a new [transformer encoder](TransformerEncoder) module.
    pub fn init<B: Backend>(&self, device: &B::Device) -> TransformerEncoder<B> {
impl<B: Backend> TransformerEncoder<B> {
/// Applies the forward pass on the input tensor.
///
/// # Shapes
///
/// - tensor: `[batch_size, seq_length, d_model]`
/// - output: `[batch_size, seq_length, d_model]`
    pub fn forward(&self, input: TransformerEncoderInput<B>) -> Tensor<B, 3> {
/// Applies the forward pass on the input tensor using autoregressive cache.
///
/// # Shapes
///
/// - tensor: `[batch_size, seq_length, d_model]`
/// - output: `[batch_size, seq_length, d_model]`
    pub fn forward_autoregressive_inference(
        &self,
        input: TransformerEncoderInput<B>,
        cache: &mut TransformerEncoderAutoregressiveCache<B>,
    ) -> Tensor<B, 3> {
/// Create an empty autoregressive cache.
    pub fn new_autoregressive_cache(&self) -> TransformerEncoderAutoregressiveCache<B> {
/// Transformer encoder layer module.
pub struct TransformerEncoderLayer<B: Backend> {
impl<B: Backend> TransformerEncoderLayer<B> {
    fn new(config: &TransformerEncoderConfig, device: &B::Device) -> Self {
    fn forward(
        &self,
        input: Tensor<B, 3>,
        mask_pad: Option<Tensor<B, 2, Bool>>,
        mask_attn: Option<Tensor<B, 3, Bool>>,
    ) -> Tensor<B, 3> {
// Multi-head attention residual path.
// Normalize.
// Multi-head attention.
// Feed forward residual path.
// Normalize.
// Feed forward.
// Main path.
// Normalize.
    fn forward_autoregressive_inference(
        &self,
        input: Tensor<B, 3>,
        mask_pad: Option<Tensor<B, 2, Bool>>,
        mask_attn: Option<Tensor<B, 3, Bool>>,
        cache: &mut TransformerEncoderLayerAutoregressiveCache<B>,
    ) -> Tensor<B, 3> {
// Multi-head attention residual path.
// Normalize.
// Multi-head attention.
// Feed forward residual path.
// Normalize.
// Feed forward.
// Main path.
// Normalize.
struct TransformerEncoderLayerAutoregressiveCache<B: Backend> {
impl<B: Backend> TransformerEncoderLayerAutoregressiveCache<B> {
    fn empty() -> Self {
/// Autoregressive cache for the [Transformer Encoder](TransformerEncoder) layer.
///
/// To be used during inference when decoding tokens.
pub struct TransformerEncoderAutoregressiveCache<B: Backend> {
impl<B: Backend> TransformerEncoderAutoregressiveCache<B> {
    fn empty(num_layers: usize) -> Self {
    type FT = FloatElem<TestBackend>;
    fn test_autoregressive_norm_last() {
    fn test_autoregressive_norm_first() {
    fn test_autoregressive(config: TransformerEncoderConfig) {
    fn display() {
/// Configuration to create a [position-wise feed-forward](PositionWiseFeedForward) layer using the [init function](PositionWiseFeedForwardConfig::init).
pub struct PositionWiseFeedForwardConfig {
/// The size of the input and output features.
/// The size of the hidden inner features.
/// The dropout rate. Default: 0.1
/// The type of function used to initialize neural network parameters
/// Applies the position-wise feed-forward network to the input tensor from the paper [Attention Is All You Need](https://arxiv.org/pdf/1706.03762v7).
///
/// # Params
///
/// - linear inner: Linear layer with `d_model` input features and `d_ff` output features.
/// - linear outer: Linear layer with `d_ff` input features and `d_model` output features.
///
/// `FFN(x) = max(0, xW1 + b1)W2 + b2`
///
/// Should be created using [PositionWiseFeedForwardConfig]
pub struct PositionWiseFeedForward<B: Backend> {
/// Linear layer with `d_model` input features and `d_ff` output features.
/// Linear layer with `d_ff` input features and `d_model` output features.
/// Dropout layer.
/// GELU activation function.
impl<B: Backend> ModuleDisplay for PositionWiseFeedForward<B> {
    fn custom_settings(&self) -> Option<DisplaySettings> {
    fn custom_content(&self, content: Content) -> Option<Content> {
impl PositionWiseFeedForwardConfig {
/// Initialize a new [position-wise feed-forward](PositionWiseFeedForward) module.
    pub fn init<B: Backend>(&self, device: &B::Device) -> PositionWiseFeedForward<B> {
impl<B: Backend> PositionWiseFeedForward<B> {
/// Applies the forward pass on the input tensor.
///
/// # Shapes
///
/// - tensor: `[batch_size, seq_length, d_model]`
/// - output: `[batch_size, seq_length, d_model]`
    pub fn forward<const D: usize>(&self, input: Tensor<B, D>) -> Tensor<B, D> {
    fn display() {
/// Configuration to create an [unfold 4d](Unfold4d) layer using the [init function](Unfold4dConfig::init).
pub struct Unfold4dConfig {
/// The size of the kernel.
/// The stride of the convolution.
/// Spacing between kernel elements.
/// The padding configuration.
/// Four-dimensional unfolding.
///
/// Should be created with [Unfold4dConfig].
pub struct Unfold4d {
/// The size of the kernel.
/// The stride of the convolution.
/// Spacing between kernel elements.
/// The padding configuration.
impl ModuleDisplay for Unfold4d {
    fn custom_settings(&self) -> Option<DisplaySettings> {
    fn custom_content(&self, content: Content) -> Option<Content> {
impl Unfold4dConfig {
/// Initializes a new [Unfold4d] module.
    pub fn init(&self) -> Unfold4d {
impl Unfold4d {
/// Applies the forward pass on the input tensor.
///
/// See [unfold4d](burn::tensor::module::unfold4d) for more information.
///
/// # Shapes
///
/// input:   `[batch_size, channels_in, height, width]`
/// returns: `[batch_size, channels_in * kernel_size_1 * kernel_size_2, number of blocks]`
    pub fn forward<B: Backend>(&self, input: Tensor<B, 4>) -> Tensor<B, 3> {
    fn display() {
/// Padding configuration for 1D operators.
pub enum PaddingConfig1d {
/// Dynamically calculates padding to ensure output size matches input size.
/// No padding applied.
/// Applies a specific amount of padding to all inputs.
impl PaddingConfig1d {
    pub(crate) fn calculate_padding_1d(
        &self,
        length: usize,
        kernel_size: usize,
        stride: usize,
    ) -> usize {
/// Padding configuration for 2D operators.
pub enum PaddingConfig2d {
/// Dynamically calculates padding to preserve input dimensions in output.
/// No padding applied.
/// Applies specified padding values to height and width dimensions.
impl PaddingConfig2d {
    pub(crate) fn calculate_padding_2d(
        &self,
        height: usize,
        width: usize,
        kernel_size: &[usize;
/// Padding configuration for 3D operators.
pub enum PaddingConfig3d {
/// Dynamically calculates padding to preserve input dimensions in output.
/// No padding applied.
/// Applies specified padding values to depth, height, and width dimensions.
impl PaddingConfig3d {
    pub(crate) fn calculate_padding_3d(
        &self,
        depth: usize,
        height: usize,
        width: usize,
        kernel_size: &[usize;
pub type B = burn_ndarray::NdArray<f32>;
/// Backend for test cases
pub type B = burn_wgpu::Wgpu;
/// Backend for test cases
pub type B = burn_cuda::Cuda;
/// Backend for test cases
pub type B = burn_rocm::Rocm;
fn should_quantize_module<M: Module<B>, const D: usize, F: Fn(&M) -> Tensor<B, D>>(
    module: M,
    scheme: QuantScheme,
    func: F,
    tolerance: Tolerance<FloatElem<B>>,
) {
fn should_quantize_transformer() {
// slightly higher abs tolerance (permissive: 1e-2)
fn should_quantize_linear_128_256() {
fn should_quantize_linear() {
// Default scheme should select supported QuantStore default
// TODO: set native if dtype is supported by the test backend
// .with_store(QuantStore::Native)
fn should_quantize_linear_weights() {
fn should_quantize_linear_blocks() {
// .with_store(QuantStore::Native)
fn should_quantize_linear_weights_blocks() {
// .with_store(QuantStore::Native)
/// Gradient Clipping provides a way to mitigate exploding gradients
pub enum GradientClippingConfig {
/// Clip the gradient by value.
/// Clip the gradient by norm.
impl GradientClippingConfig {
/// Initialize the gradient clipping.
///
/// # Returns
///
/// The gradient clipping.
    pub fn init(&self) -> GradientClipping {
/// Gradient Clipping provides a way to mitigate exploding gradients
/// by clipping every component of the gradient by value or by norm during
/// backpropagation.
pub enum GradientClipping {
/// Clip the gradient by value.
/// Clip the gradient by norm.
impl GradientClipping {
/// Clip the gradient.
///
/// # Arguments
///
/// * `grad` - The gradient to clip.
///
/// # Returns
///
/// The clipped gradient.
    pub fn clip_gradient<B: Backend, const D: usize>(&self, grad: Tensor<B, D>) -> Tensor<B, D> {
    fn clip_by_value<B: Backend, const D: usize>(
        &self,
        grad: Tensor<B, D>,
        threshold: f32,
    ) -> Tensor<B, D> {
    fn clip_by_norm<B: Backend, const D: usize>(
        &self,
        grad: Tensor<B, D>,
        threshold: f32,
    ) -> Tensor<B, D> {
// avoid div by zero
    fn l2_norm<B: Backend, const D: usize>(tensor: Tensor<B, D>) -> Tensor<B, 1> {
    fn test_clip_by_value() {
    fn test_clip_by_norm() {
    fn test_clip_by_norm_no_clipping() {
//! Burn optimizers.
/// Optimizer module.
/// Gradient clipping module.
/// Learning rate scheduler module.
/// Type alias for the learning rate.
///
/// LearningRate also implements [learning rate scheduler](crate::lr_scheduler::LrScheduler) so it
/// can be used for constant learning rate.
pub type LearningRate = f64;
/// Backend for test cases
pub type TestBackend = burn_ndarray::NdArray<f32>;
/// Backend for test cases
pub type TestBackend = burn_tch::LibTorch<f32>;
/// Backend for test cases
pub type TestBackend = burn_wgpu::Wgpu;
/// Backend for test cases
pub type TestBackend = burn_cuda::Cuda;
/// Backend for test cases
pub type TestBackend = burn_rocm::Rocm;
/// Backend for autodiff test cases
pub type TestAutodiffBackend = burn_autodiff::Autodiff<TestBackend>;
/// Learning rate scheduler defines how the learning rate will evolve during training.
pub trait LrScheduler: Clone + Send + Sync {
/// Scheduler associative type to be used when saving and loading the state.
    type Record<B: Backend>: Record<B>;
/// Perform the scheduler step, potentially updating its state, and returning the effective
/// learning rate.
    fn step(&mut self) -> LearningRate;
/// Get the current state of the scheduler as a [record](Record).
    fn to_record<B: Backend>(&self) -> Self::Record<B>;
/// Load the state of the scheduler as a [record](Record).
    fn load_record<B: Backend>(self, record: Self::Record<B>) -> Self;
// A small tolerance for learning rate comparisons. Depending on how learning rates are
// computed, floating-point arithmetic error might exceed f64::EPSILON, so a larger value is
// used here.
    const LOOSE_EPSILON: LearningRate = 1e-10;
    pub fn check_lr_sequence<I, S>(mut scheduler: S, expected_lrs: I)
    where
        I: IntoIterator<Item = LearningRate>,
        S: LrScheduler,
    {
// save_at_step is the number of steps to run the scheduler before saving and loading back its
// state.
    pub fn check_save_load<S>(mut scheduler: S, save_at_step: usize)
    where
        S: Clone + LrScheduler,
    {
// Consume some steps before saving and loading back
// Validate that the scheduler resumes from where it left off.
// Check if two schedulers produce the same learning rate sequences over the specified number of
// steps.
    pub fn compare_steps<S: LrScheduler>(a: &mut S, b: &mut S, num_steps: usize) {
/// Compose multiple [learning rate schedulers](LrScheduler) together.
pub struct ComposedLrSchedulerConfig {
/// Compose multiple [learning rate schedulers](LrScheduler) together.
pub struct ComposedLrScheduler {
/// Defines how the learning rates generated by the schedulers are combined.
pub enum SchedulerReduction {
/// All learning rates are averaged.
/// All learning rates are summed.
/// All learning rates are multiplied.
impl ComposedLrSchedulerConfig {
/// Initialize the learning rate scheduler.
    pub fn init(&self) -> Result<ComposedLrScheduler, String> {
/// Appends a [linear scheduler](LinearLrScheduler).
    pub fn linear(mut self, config: LinearLrSchedulerConfig) -> Self {
/// Appends a [cosine scheduler](ComposedLrSchedulerConfig).
    pub fn cosine(mut self, config: CosineAnnealingLrSchedulerConfig) -> Self {
/// Appends an [exponential scheduler](ExponentialLrScheduler).
    pub fn exponential(mut self, config: ExponentialLrSchedulerConfig) -> Self {
/// Appends a [noam scheduler](NoamLrScheduler).
    pub fn noam(mut self, config: NoamLrSchedulerConfig) -> Self {
enum LrSchedulerConfig {
enum LrSchedulerItem {
/// Record item for the [componsed learning rate scheduler](ComposedLrScheduler).
pub enum LrSchedulerRecord<B: Backend> {
/// The linear variant.
/// The cosine variant.
/// The exponential variant.
/// The noam variant.
/// Records for the [componsed learning rate scheduler](ComposedLrScheduler).
pub struct ComposedLrSchedulerRecord<B: Backend> {
impl LrScheduler for ComposedLrScheduler {
    type Record<B: Backend> = ComposedLrSchedulerRecord<B>;
    fn step(&mut self) -> LearningRate {
    fn to_record<B: Backend>(&self) -> Self::Record<B> {
    fn load_record<B: Backend>(mut self, record: Self::Record<B>) -> Self {
/// Constant learning rate implementing [learning rate scheduler](LrScheduler).
///
/// # Notes
///
/// You can also use [learning rate](LearningRate) which the same effect.
pub struct ConstantLr {
impl From<LearningRate> for ConstantLr {
    fn from(lr: LearningRate) -> Self {
impl LrScheduler for ConstantLr {
    type Record<B: Backend> = ();
    fn step(&mut self) -> LearningRate {
    fn to_record<B: Backend>(&self) -> Self::Record<B> {
    fn load_record<B: Backend>(self, _record: Self::Record<B>) -> Self {
impl LrScheduler for LearningRate {
    type Record<B: Backend> = ();
    fn step(&mut self) -> LearningRate {
    fn to_record<B: Backend>(&self) -> Self::Record<B> {
    fn load_record<B: Backend>(self, _record: Self::Record<B>) -> Self {
/// The configuration for creating a [Cosine Annealing learning rate scheduler with warm
/// restarts](CosineAnnealingLrScheduler).
///
/// This scheduler returns the learning rate `initial_lr` at the first step, then changes it by
/// following a cosine function. After `num_iters` iterations, the learning rate is reset to
/// `initial_lr`.
pub struct CosineAnnealingLrSchedulerConfig {
// The initial learning rate.
// The final learning rate.
// The number of iterations between two restarts. The two restart iterations themselves are not
// included.
impl CosineAnnealingLrSchedulerConfig {
/// Initializes a [Cosine learning rate scheduler](CosineAnnealingLrScheduler).
///
/// # Errors
///
/// An error will be returned if any of the following conditions is true:
///
/// * `initial_lr` is out of range (0.0, 1.0]
/// * `min_lr` is out of range [0.0, `initial_lr`]
/// * `num_iters` is 0
    pub fn init(&self) -> Result<CosineAnnealingLrScheduler, String> {
/// A Cosine Annealing learning rate scheduler.
///
/// This scheduler is described in [SGDR: Stochastic Gradient Descent with Warm
/// Restarts](https://arxiv.org/abs/1608.03983). See [CosineAnnealingLrSchedulerConfig] for more
/// information.
pub struct CosineAnnealingLrScheduler {
impl LrScheduler for CosineAnnealingLrScheduler {
    type Record<B: Backend> = usize;
    fn step(&mut self) -> LearningRate {
// Make current_iter overflow from usize::MAX to 0 to get the initial learning rate on the
// first call. We could've used i64 with an initial value -1, but keeping it in usize saves
// us from some type casting here.
    fn to_record<B: Backend>(&self) -> Self::Record<B> {
    fn load_record<B: Backend>(mut self, record: Self::Record<B>) -> Self {
    fn config_initial_lr_too_low() {
    fn config_initial_lr_too_high() {
    fn config_min_lr_too_low() {
    fn config_min_lr_too_high() {
    fn config_num_iters_too_low() {
    fn test_lr_change() {
        const INITIAL_LR: LearningRate = 0.5;
        const MIN_LR: LearningRate = 0.1;
// cos(0)
// cos(PI/2)
// cos(PI)
// restart
    fn test_save_and_load() {
        const NUM_ITERS: usize = 9;
/// The configuration for creating an [exponential learning rate scheduler](ExponentialLrScheduler).
///
/// This scheduler returns the learning rate `initial_lr` at the first step, then multiplies it by
/// a constant `gamma` at every iteration. At any iteration `i` (which starts from 0), the learning
/// rate is given by `initial_lr * gamma^i`.
pub struct ExponentialLrSchedulerConfig {
// The initial learning rate.
// The constant that the learning rate is multiplied by on each iteration.
impl ExponentialLrSchedulerConfig {
/// Initializes a [exponential learning rate scheduler](ExponentialLrScheduler).
///
/// # Errors
///
/// An error will be returned if any of the following conditions is true:
///
/// * `initial_lr` is out of range (0.0, 1.0]
/// * `gamma` is out of range (0.0, 1.0]
    pub fn init(&self) -> Result<ExponentialLrScheduler, String> {
// Such an initial value eliminates the need for special-case handling of the first
// learning rate.
/// A exponential learning rate scheduler.
///
/// See [ExponentialLrSchedulerConfig] for more information.
pub struct ExponentialLrScheduler {
// The previous iteration's learning rate.
// The constant that the learning rate is multiplied by on each iteration.
impl LrScheduler for ExponentialLrScheduler {
    type Record<B: Backend> = LearningRate;
    fn step(&mut self) -> LearningRate {
    fn to_record<B: Backend>(&self) -> Self::Record<B> {
    fn load_record<B: Backend>(mut self, record: Self::Record<B>) -> Self {
    fn config_initial_lr_too_low() {
    fn config_initial_lr_too_high() {
    fn config_gamma_too_low() {
    fn config_gamma_too_high() {
    fn test_lr_change() {
    fn test_save_and_load() {
/// The configuration for creating a [linear learning rate scheduler](LinearLrScheduler).
///
/// This scheduler returns the learning rate `initial_lr` at the first step, then changes it by a
/// constant amount on each iteration until reaching a final learning rate `final_lr`. The
/// `num_iters` parameter controls how many iterations are needed to go from `initial_lr` to
/// `final_lr`.
pub struct LinearLrSchedulerConfig {
// The initial learning rate.
// The final learning rate.
// The number of iterations before reaching the final learning rate.
impl LinearLrSchedulerConfig {
/// Initializes a [linear learning rate scheduler](LinearLrScheduler).
///
/// # Errors
///
/// An error will be returned if any of the following conditions is true:
///
/// * `initial_lr` is out of range (0.0, 1.0]
/// * `final_lr` is out of range [0.0, 1.0]
/// * `num_iters` is 0
    pub fn init(&self) -> Result<LinearLrScheduler, String> {
/// A linear learning rate scheduler.
///
/// See [LinearLrSchedulerConfig] for more information.
pub struct LinearLrScheduler {
// The final learning rate after the linear changing process stops.
// The amount that the learning rate changes by on each iteration.
// The number of iterations left before reaching the final learning rate.
impl LrScheduler for LinearLrScheduler {
    type Record<B: Backend> = usize;
    fn step(&mut self) -> LearningRate {
    fn to_record<B: Backend>(&self) -> Self::Record<B> {
    fn load_record<B: Backend>(mut self, record: Self::Record<B>) -> Self {
    fn config_initial_lr_too_low() {
    fn config_initial_lr_too_high() {
    fn config_final_lr_too_low() {
    fn config_final_lr_too_high() {
    fn config_num_iters_too_low() {
    fn test_lr_decreasing() {
    fn test_lr_increasing() {
    fn test_lr_unchanging() {
    fn test_save_and_load() {
        const NUM_ITERS: usize = 6;
/// Constant learning rate scheduler
/// Composed learning rate scheduler
/// Linear learning rate scheduler
/// Noam learning rate scheduler
/// Exponential learning rate scheduler
/// Cosine learning rate scheduler
/// Step learning rate scheduler
/// Configuration to create a [noam](NoamLrScheduler) learning rate scheduler.
pub struct NoamLrSchedulerConfig {
/// The overall scale factor for the learning rate decay.
/// The number of steps before the exponential decay stats.
/// The size of the model.
/// Noam learning rate scheduler as described in [Attention Is All You Need](https://arxiv.org/abs/1706.03762).
pub struct NoamLrScheduler {
impl NoamLrSchedulerConfig {
/// Initialize a new [noam](NoamLrScheduler) learning rate scheduler.
///
/// # Errors
///
/// An error will be returned if any of the following conditions is true:
///
/// * `warmup_steps` is 0
/// * `model_size` is 0
    pub fn init(&self) -> Result<NoamLrScheduler, String> {
impl LrScheduler for NoamLrScheduler {
    type Record<B: Backend> = usize;
    fn step(&mut self) -> LearningRate {
    fn to_record<B: Backend>(&self) -> Self::Record<B> {
    fn load_record<B: Backend>(mut self, record: Self::Record<B>) -> Self {
    fn test_config_warmup_steps_invalid() {
    fn test_config_warmup_steps_valid() {
    fn test_config_model_size_invalid() {
    fn test_config_model_size_valid() {
    fn test_function_increase_and_decrease() {
/// The configuration for create a [step learning rate scheduler](StepLrScheduler).
///
/// This scheduler returns the learning rate `initial_lr` from the start, and keeps doing so until
/// the same value has been given for `step_size` times. Then it multiplies the learning rate by
/// `gamma` before repeating the process.
///
/// Gamma values out of range (0.0, 1.0) and non-positive initial learning rates are acceptable, but
/// a warning log will be output for such a value in case of mistyping.
///
/// ## Notes
///
/// The [step](StepLrScheduler::step) method of the scheduler panics if it is called more than
/// `i32::MAX + 1` times.
pub struct StepLrSchedulerConfig {
// The learning rate at the initial step.
// The number of iterations over which the learning rate remains unchanged before the next
// update.
/// The factor by which the learning rate is multiplied with each update. Default: 0.1.
impl StepLrSchedulerConfig {
/// Initializes a [step learning rate scheduler](StepLrScheduler).
///
/// # Errors
///
/// An error will be returned if `step_size` is 0.
    pub fn init(&self) -> Result<StepLrScheduler, String> {
// Atypical values of `initial_lr` and `gamma` are not rejected because they might be useful
// in some cases like debugging (e.g., https://datascience.stackexchange.com/q/89518).
/// Step learning rate scheduler.
pub struct StepLrScheduler {
// The index of the current iteration.
// `i32` is used for avoiding truncating the exponent when taking powers of `gamma`.
impl LrScheduler for StepLrScheduler {
    type Record<B: Backend> = i32;
    fn step(&mut self) -> LearningRate {
// Type casting below causes no truncation, as all the values fall within the ranges.
    fn to_record<B: Backend>(&self) -> Self::Record<B> {
    fn load_record<B: Backend>(mut self, record: Self::Record<B>) -> Self {
// Warning logs for initial LR and gamma are not tested because there seems no straightforward
// way to do it.
//
// Creating a mock logger that collects logs into `String` for later examination seems a possible
// solution, but unit tests run in the same process in parallel, where the single logger would
// be shared by multiple tests, so logs from different tests would be mixed up with no easy way
// to separate them.
// Using "--test-threads=1" could prevent mixup, but whether the ability to test logging is
// worth the slowdown would be a question. Also, using a primitive provided by `std` to
// synchronize the logger across tests is not an option since we need to support `no-std`.
// Maybe the mocking approach can be reconsidered after we are given an option to run tests in
// separate processes like what the issue below is proposing:
//     https://github.com/rust-lang/rust/issues/47506
//
// As a side note, a helper crate exists for the exact purpose:
//     https://crates.io/crates/testing_logger
// but the crate has been unmaintained and using it would introduce another dependency.
    fn test_config_step_size_zero() {
    fn test_config_step_size_nonzero() {
    fn test_config_default_gamma() {
        const INIT_LR: LearningRate = 0.4;
        const STEP_SIZE: usize = 2;
    fn test_lr_decreasing() {
    fn test_lr_increasing() {
    fn test_lr_unchanging() {
    fn test_save_and_load() {
        const STEP_SIZE: usize = 10;
// It's too time consuming to actually run a scheduler `i32::MAX` steps, so an approach that
// depends on private fields is used to implement the test.
    fn test_number_of_calls_within_limit() {
// Create a scheduler that has already run `i32::MAX` steps
    fn test_number_of_calls_over_limit() {
// Create a scheduler that has already run `i32::MAX` steps
/// AdaGrad configuration.
pub struct AdaGradConfig {
/// [Weight decay](WeightDecayConfig) config.
/// [Gradient Clipping](GradientClippingConfig) config.
/// AdaGrad optimizer
pub struct AdaGrad {
/// AdaGrad state.
pub struct AdaGradState<B: Backend, const D: usize> {
impl<B: Backend> SimpleOptimizer<B> for AdaGrad {
    type State<const D: usize> = AdaGradState<B, D>;
    fn step<const D: usize>(
        &self,
        lr: LearningRate,
        tensor: Tensor<B, D>,
        mut grad: Tensor<B, D>,
        state: Option<Self::State<D>>,
    ) -> (Tensor<B, D>, Option<Self::State<D>>) {
    fn to_device<const D: usize>(mut state: Self::State<D>, device: &Device<B>) -> Self::State<D> {
impl AdaGradConfig {
/// Initialize AdaGrad optimizer.
///
/// # Returns
///
/// Returns an optimizer that can be used to optimize a module.
    pub fn init<B: AutodiffBackend, M: AutodiffModule<B>>(
        &self,
    ) -> OptimizerAdaptor<AdaGrad, M, B> {
/// Learning rate decay state (also includes sum state).
pub struct LrDecayState<B: Backend, const D: usize> {
struct LrDecay {
impl LrDecay {
    pub fn transform<B: Backend, const D: usize>(
        &self,
        grad: Tensor<B, D>,
        lr: LearningRate,
        lr_decay_state: Option<LrDecayState<B, D>>,
    ) -> (Tensor<B, D>, LrDecayState<B, D>) {
impl<B: Backend, const D: usize> LrDecayState<B, D> {
/// Move state to device.
///
/// # Arguments
///
/// * `device` - Device to move state to.
///
/// # Returns
///
/// Returns state moved to device.
    pub fn to_device(mut self, device: &B::Device) -> Self {
    const LEARNING_RATE: LearningRate = 0.01;
    fn test_adagrad_optimizer_save_load_state() {
    fn test_adagrad_optimizer_with_numbers() {
        type FT = FloatElem<TestAutodiffBackend>;
    fn given_linear_layer(weight: TensorData, bias: TensorData) -> Linear<TestAutodiffBackend> {
    fn create_adagrad()
    -> OptimizerAdaptor<AdaGrad, Linear<TestAutodiffBackend>, TestAutodiffBackend> {
/// Adam configuration.
pub struct AdamConfig {
/// Parameter for Adam.
/// Parameter for Adam.
/// A value required for numerical stability.
/// [Weight decay](WeightDecayConfig) config.
/// [Gradient Clipping](GradientClippingConfig) config.
/// Adam optimizer as described in the paper [Adam: A Method for Stochastic Optimization](https://arxiv.org/pdf/1412.6980.pdf).
pub struct Adam {
/// Adam state.
pub struct AdamState<B: Backend, const D: usize> {
/// The current adaptive momentum.
impl<B: Backend> SimpleOptimizer<B> for Adam {
    type State<const D: usize> = AdamState<B, D>;
    fn step<const D: usize>(
        &self,
        lr: LearningRate,
        tensor: Tensor<B, D>,
        mut grad: Tensor<B, D>,
        state: Option<Self::State<D>>,
    ) -> (Tensor<B, D>, Option<Self::State<D>>) {
    fn to_device<const D: usize>(mut state: Self::State<D>, device: &Device<B>) -> Self::State<D> {
impl AdamConfig {
/// Initialize Adam optimizer.
///
/// # Returns
///
/// Returns an optimizer that can be used to optimize a module.
    pub fn init<B: AutodiffBackend, M: AutodiffModule<B>>(&self) -> OptimizerAdaptor<Adam, M, B> {
/// Adaptive momentum state.
pub struct AdaptiveMomentumState<B: Backend, const D: usize> {
/// The number of iterations aggregated.
/// The first order momentum.
/// The second order momentum.
struct AdaptiveMomentum {
impl AdaptiveMomentum {
    pub fn transform<B: Backend, const D: usize>(
        &self,
        grad: Tensor<B, D>,
        momentum_state: Option<AdaptiveMomentumState<B, D>>,
    ) -> (Tensor<B, D>, AdaptiveMomentumState<B, D>) {
impl<B: Backend, const D: usize> AdaptiveMomentumState<B, D> {
/// Move state to device.
///
/// # Arguments
///
/// * `device` - Device to move state to.
///
/// # Returns
///
/// Returns state moved to device.
    pub fn to_device(mut self, device: &B::Device) -> Self {
    const LEARNING_RATE: LearningRate = 0.01;
    fn test_adam_optimizer_save_load_state() {
    fn test_adam_optimizer_with_numbers() {
        type FT = FloatElem<TestAutodiffBackend>;
    fn test_adam_optimizer_no_nan() {
    fn given_linear_layer(weight: TensorData, bias: TensorData) -> Linear<TestAutodiffBackend> {
    fn create_adam() -> OptimizerAdaptor<Adam, Linear<TestAutodiffBackend>, TestAutodiffBackend> {
/// [`AdamW`] Configuration.
pub struct AdamWConfig {
/// Parameter for AdamW.
/// Parameter for AdamW.
/// A value required for numerical stability.
/// Weight decay config.
/// Cautious weight decay config.
///
/// See: <https://arxiv.org/abs/2510.12402>
/// [Gradient Clipping](GradientClippingConfig) config.
/// AdamW optimizer.
///
/// See:
/// - [Decoupled Weight Decay Regularization, Loshchilov and Hutter, 2019](https://arxiv.org/abs/1711.05101).
/// - [Cautious Weight Decay, 2025](https://arxiv.org/abs/2510.12402)
///
/// Configured by [`AdamWConfig`].
pub struct AdamW {
/// AdamW state.
pub struct AdamWState<B: Backend, const D: usize> {
/// Th current adaptive momentum state.
impl<B: Backend> SimpleOptimizer<B> for AdamW {
    type State<const D: usize> = AdamWState<B, D>;
/// A single optimization step for any tensor that represents the parameters of a model.
    fn step<const D: usize>(
        &self,
        // Learning rate.
        lr: LearningRate,
        // Any tensor that represents the parameters of a model.
        tensor: Tensor<B, D>,
        // Gradient of the loss w.r.t. the parameters.
        grad: Tensor<B, D>,
        // State of the optimizer.
        state: Option<Self::State<D>>,
    ) -> (Tensor<B, D>, Option<Self::State<D>>) {
// Cautious weight decay.
// See: https://arxiv.org/abs/2510.12402
// Zero out the decay where the decay is counter to the update direction.
    fn to_device<const D: usize>(mut state: Self::State<D>, device: &Device<B>) -> Self::State<D> {
impl AdamWConfig {
/// Initialize AdamW optimizer.
///
/// # Returns
///
/// Returns an optimizer that can be used to optimize a module.
    pub fn init<B: AutodiffBackend, M: AutodiffModule<B>>(&self) -> OptimizerAdaptor<AdamW, M, B> {
struct AdaptiveMomentumW {
impl AdaptiveMomentumW {
    pub fn transform<B: Backend, const D: usize>(
        &self,
        grad: Tensor<B, D>,
        state: Option<AdaptiveMomentumState<B, D>>,
    ) -> (Tensor<B, D>, AdaptiveMomentumState<B, D>) {
// Update first moment estimate.
// Update second moment estimate.
// Update time.
// Initialize first moment estimate.
// Initialize second moment estimate.
// Compute bias-corrected first and second moment estimates.
// Compute update delta. This still needs to be scaled by the learning rate.
    type FT = FloatElem<TestAutodiffBackend>;
    const LEARNING_RATE: LearningRate = 0.01;
    fn test_adamw_optimizer_save_load_state() {
    fn test_adamw_optimizer_with_numbers() {
    fn test_adamw_optimizer_with_numbers_cautious() {
    fn test_adam_optimizer_no_nan() {
    fn given_linear_layer(weight: TensorData, bias: TensorData) -> Linear<TestAutodiffBackend> {
    fn create_adamw() -> OptimizerAdaptor<AdamW, Linear<TestAutodiffBackend>, TestAutodiffBackend> {
/// General trait to optimize [module](AutodiffModule).
pub trait Optimizer<M, B>: Send + Clone
where
    M: AutodiffModule<B>,
    B: AutodiffBackend,
{
/// Optimizer associative type to be used when saving and loading the state.
    type Record: Record<B>;
/// Perform the optimizer step using the given learning rate and gradients.
/// The updated module is returned.
    fn step(&mut self, lr: LearningRate, module: M, grads: GradientsParams) -> M;
/// Get the current state of the optimizer as a [record](Record).
    fn to_record(&self) -> Self::Record;
/// Load the state of the optimizer as a [record](Record).
    fn load_record(self, record: Self::Record) -> Self;
/// Configuration to create [weight decay](WeightDecay).
pub struct WeightDecayConfig {
/// L2 penalty.
/// State of [weight decay](WeightDecay).
pub struct WeightDecayState<B: Backend, const D: usize> {
/// Weight decay implementation that transforms gradients.
pub struct WeightDecay {
impl WeightDecay {
/// Creates a new [weight decay](WeightDecay) from a [config](WeightDecayConfig).
    pub fn new(config: &WeightDecayConfig) -> Self {
/// Transforms a gradient.
///
/// # Arguments
///
/// * `grad` - Gradient to transform.
/// * `tensor` - Tensor param of the last iteration.
///
/// # Returns
///
/// * `grad` - Transformed gradient.
    pub fn transform<B: Backend, const D: usize>(
        &self,
        grad: Tensor<B, D>,
        tensor: Tensor<B, D>,
    ) -> Tensor<B, D> {
impl<B: Backend, const D: usize> WeightDecayState<B, D> {
/// Moves the state to a device.
///
/// # Arguments
///
/// * `device` - Device to move the state to.
///
/// # Returns
///
/// * `self` - Moved state.
    pub fn to_device(mut self, device: &B::Device) -> Self {
/// Accumulate gradients into a single [Gradients](AutodiffBackend::Gradients) object.
pub struct GradientsAccumulator<M> {
impl<M> Default for GradientsAccumulator<M> {
    fn default() -> Self {
impl<M> GradientsAccumulator<M> {
/// Create a new gradients accumulator.
    pub fn new() -> Self {
impl<M> GradientsAccumulator<M> {
/// Accumulate the given gradients for each parameter in the given module.
    pub fn accumulate<B: AutodiffBackend>(&mut self, module: &M, grads: GradientsParams)
    where
        M: AutodiffModule<B>,
    {
/// Return the accumulated gradients and reset the accumulator state.
    pub fn grads(&mut self) -> GradientsParams {
struct ModuleGradsAccumulator<'a, M> {
impl<B: AutodiffBackend, M: AutodiffModule<B>> ModuleVisitor<B> for ModuleGradsAccumulator<'_, M> {
    fn visit_float<const D: usize>(&mut self, param: &Param<Tensor<B, D>>) {
    fn test_accumulate_gradients_one_step() {
    fn test_accumulate_gradients_two_steps() {
    fn layer<B: Backend>(device: &B::Device) -> Linear<B> {
    fn random_tensor<B: Backend>(device: &B::Device) -> Tensor<B, 2> {
/// Data type that contains gradients for parameters.
pub struct GradientsParams {
impl GradientsParams {
/// Creates a new [GradientsParams](GradientsParams).
    pub fn new() -> Self {
/// Extract each tensor gradients for the given [module](AutodiffModule).
///
/// Note: This consumes the gradients. See ['from_module'] to extract gradients only for
///  a specific module.
    pub fn from_grads<B: AutodiffBackend, M: AutodiffModule<B>>(
        grads: B::Gradients,
        module: &M,
    ) -> Self {
/// Extract each tensor gradients for the given [module](AutodiffModule).
    pub fn from_module<B: AutodiffBackend, M: AutodiffModule<B>>(
        grads: &mut B::Gradients,
        module: &M,
    ) -> Self {
/// Extract tensor gradients for the given [module](AutodiffModule) and given parameters.
    pub fn from_params<B: AutodiffBackend, M: AutodiffModule<B>>(
        grads: &mut B::Gradients,
        module: &M,
        params: &[ParamId],
    ) -> Self {
/// Get the gradients for the given [parameter id](ParamId).
///
/// # Notes
///
/// You should use [remove](GradientsParams::remove) if you want to get the gradients
/// only one time.
    pub fn get<B, const D: usize>(&self, id: ParamId) -> Option<Tensor<B, D>>
    where
        B: Backend,
    {
/// Remove the gradients for the given [parameter id](ParamId).
    pub fn remove<B, const D: usize>(&mut self, id: ParamId) -> Option<Tensor<B, D>>
    where
        B: Backend,
    {
/// Register a gradients tensor for the given [parameter id](ParamId).
///
/// # Notes
///
/// If a tensor is already registered for the given [parameter id](ParamId), it will be replaced.
    pub fn register<B, const D: usize>(&mut self, id: ParamId, value: Tensor<B, D>)
    where
        B: Backend,
    {
/// The number of gradients tensors registered.
    pub fn len(&self) -> usize {
/// If any tensor is contained.
    pub fn is_empty(&self) -> bool {
/// Change the device of each tensor gradients registered for the given [module](AutodiffModule).
    pub fn to_device<B: AutodiffBackend, M: AutodiffModule<B>>(
        mut self,
        device: &B::Device,
        module: &M,
    ) -> Self {
/// Syncs the gradient params with the other peers in the collective.
    pub fn all_reduce<B: Backend>(
        mut self,
        peer_id: PeerId,
        op: ReduceOperation,
    ) -> Result<Self, CollectiveError> {
// This is crucial, since the all-reduce operations need to happen in the same order for the same parameters on all nodes!
    fn test_convert_grads() {
    fn layer<B: Backend>(device: &B::Device) -> Linear<B> {
    fn random_tensor<B: Backend>(device: &B::Device) -> Tensor<B, 2> {
/// Weight decay module for optimizers.
/// Momentum module for optimizers.
/// Configuration to create [momentum](Momentum).
pub struct MomentumConfig {
/// Momentum factor
/// Dampening factor.
/// Enables Nesterov momentum, see [On the importance of initialization and
/// momentum in deep learning](http://www.cs.toronto.edu/~hinton/absps/momentum.pdf).
/// State of [momentum](Momentum).
pub struct MomentumState<B: Backend, const D: usize> {
/// Momentum implementation that transforms gradients.
pub struct Momentum<B: Backend> {
impl<B: Backend> Momentum<B> {
/// Creates a new [momentum](Momentum) from a [config](MomentumConfig).
    pub fn new(config: &MomentumConfig) -> Self {
/// Transforms a gradient.
///
/// # Arguments
///
/// * `grad` - Gradient to transform.
/// * `state` - State of the optimizer.
///
/// # Returns
///
/// * `grad` - Transformed gradient.
/// * `state` - State of the optimizer.
    pub fn transform<const D: usize>(
        &self,
        grad: Tensor<B, D>,
        state: Option<MomentumState<B, D>>,
    ) -> (Tensor<B, D>, MomentumState<B, D>) {
impl<B: Backend, const D: usize> MomentumState<B, D> {
/// Moves the state to a device.
///
/// # Arguments
///
/// * `device` - Device to move the state to.
///
/// # Returns
///
/// * `self` - Moved state.
    pub fn to_device(mut self, device: &B::Device) -> Self {
/// Configuration to create the [RmsProp](RmsProp) optimizer.
pub struct RmsPropConfig {
/// Smoothing constant.
/// momentum for RmsProp.
/// A value required for numerical stability.
/// if True, compute the centered RmsProp, the gradient is normalized by an estimation of its variance
/// [Weight decay](WeightDecayConfig) config.
/// [Gradient Clipping](GradientClippingConfig) config.
impl RmsPropConfig {
/// Initialize RmsProp optimizer.
///
/// # Returns
///
/// Returns an optimizer that can be used to optimize a module.
    pub fn init<B: AutodiffBackend, M: AutodiffModule<B>>(
        &self,
    ) -> OptimizerAdaptor<RmsProp, M, B> {
/// Optimizer that implements stochastic gradient descent with momentum.
/// The optimizer can be configured with [RmsPropConfig](RmsPropConfig).
pub struct RmsProp {
// epsilon: f32,
// momentum: Option<Momentum<B>>,
impl<B: Backend> SimpleOptimizer<B> for RmsProp {
    type State<const D: usize> = RmsPropState<B, D>;
    fn step<const D: usize>(
        &self,
        lr: LearningRate,
        tensor: Tensor<B, D>,
        mut grad: Tensor<B, D>,
        state: Option<Self::State<D>>,
    ) -> (Tensor<B, D>, Option<Self::State<D>>) {
// fetch state for params
// weight_decay transform
// square_avg transform
// centered transform
// momentum transform
// transition state
// tensor param transform
    fn to_device<const D: usize>(mut state: Self::State<D>, device: &Device<B>) -> Self::State<D> {
/// State of [RmsProp](RmsProp)
pub struct RmsPropState<B: Backend, const D: usize> {
/// Current squared average state.
/// Current centered state
/// Current gradient momentum, if any.
/// [SquareAvgState](SquareAvgState) is to store and pass optimizer step params.
pub struct SquareAvgState<B: Backend, const D: usize> {
/// Current squared average.
impl<B: Backend, const D: usize> SquareAvgState<B, D> {
/// transform [SquareAvgState] to the next step
    fn transform(alpha: f32, grad: Tensor<B, D>, state: Option<Self>) -> (Tensor<B, D>, Self) {
/// Moves the state to a device.
///
/// # Arguments
///
/// * `device` - Device to move the state to.
///
/// # Returns
///
/// * `self` - Moved state.
    pub fn to_device(mut self, device: &B::Device) -> Self {
/// [CenteredState](CenteredState) is to store and pass optimizer step params.
pub struct CenteredState<B: Backend, const D: usize> {
/// The averaged gradient to calculate the centered gradient, if available.
/// The current average value.
impl<B: Backend, const D: usize> CenteredState<B, D> {
/// transform [CenteredState] to the next step
    fn transform(
        alpha: f32,
        centered: bool,
        grad: Tensor<B, D>,
        square_avg_state: SquareAvgState<B, D>,
        centered_state: Option<Self>,
    ) -> (Tensor<B, D>, SquareAvgState<B, D>, Self) {
/// Moves the state to a device.
///
/// # Arguments
///
/// * `device` - Device to move the state to.
///
/// # Returns
///
/// * `self` - Moved state.
    pub fn to_device(mut self, device: &B::Device) -> Self {
/// [RmsPropMomentum](RmsPropMomentum) is to store config status for optimizer.
/// (, which is stored in [optimizer](RmsProp) itself and not passed in during `step()` calculation)
pub struct RmsPropMomentum {
impl RmsPropMomentum {
/// transform [grad](Tensor) and [RmsPropMomentumState] to the next step
    fn transform<B: Backend, const D: usize>(
        &self,
        grad: Tensor<B, D>,
        centered_state: CenteredState<B, D>,
        momentum_state: Option<RmsPropMomentumState<B, D>>,
    ) -> (
        Tensor<B, D>,
        CenteredState<B, D>,
        Option<RmsPropMomentumState<B, D>>,
    ) {
/// [RmsPropMomentumState](RmsPropMomentumState) is to store and pass optimizer step params.
pub struct RmsPropMomentumState<B: Backend, const D: usize> {
impl<B: Backend, const D: usize> RmsPropMomentumState<B, D> {
/// Moves the state to a device.
///
/// # Arguments
///
/// * `device` - Device to move the state to.
///
/// # Returns
///
/// * `self` - Moved state.
    pub fn to_device(mut self, device: &B::Device) -> Self {
    type FT = FloatElem<TestAutodiffBackend>;
    const LEARNING_RATE: LearningRate = 0.01;
    fn test_rmsprop_optimizer_save_load_state() {
/// used for test differences and debug
    fn test_rmsprop_optimizer_with_numbers_basic() {
// println!("linear is {:?}", linear);
// println!("linear is {:?}", linear);
// println!("linear is {:?}", linear);
// println!("\nweight_updated\n{:?}", weight_updated);
// println!("\nbias_updated\n{:?}", bias_updated);
    fn test_rmsprop_optimizer_with_numbers() {
// println!("\nweight_updated\n{:?}", weight_updated);
// println!("\nbias_updated\n{:?}", bias_updated);
    fn given_linear_layer(weight: TensorData, bias: TensorData) -> Linear<TestAutodiffBackend> {
    fn create_random_tensor() -> Tensor<TestAutodiffBackend, 2> {
    fn create_rmsprop()
    -> OptimizerAdaptor<RmsProp, Linear<TestAutodiffBackend>, TestAutodiffBackend> {
/// Configuration to create the [Sgd](Sgd) optimizer.
pub struct SgdConfig {
/// [Weight decay](WeightDecayConfig) config.
/// [Momentum](MomentumConfig) config.
/// [Gradient Clipping](GradientClippingConfig) config.
/// Optimizer that implements stochastic gradient descent with momentum.
///
/// The optimizer can be configured with [SgdConfig](SgdConfig).
pub struct Sgd<B: Backend> {
/// State of [Sgd](Sgd).
pub struct SgdState<B: Backend, const D: usize> {
/// The current state of the momentum (if any).
impl SgdConfig {
/// Creates a new [SgdConfig](SgdConfig) with default values.
    pub fn init<B: AutodiffBackend, M: AutodiffModule<B>>(
        &self,
    ) -> OptimizerAdaptor<Sgd<B::InnerBackend>, M, B> {
impl<B: Backend> SimpleOptimizer<B> for Sgd<B> {
    type State<const D: usize> = SgdState<B, D>;
    fn step<const D: usize>(
        &self,
        lr: LearningRate,
        tensor: Tensor<B, D>,
        mut grad: Tensor<B, D>,
        state: Option<Self::State<D>>,
    ) -> (Tensor<B, D>, Option<Self::State<D>>) {
    fn to_device<const D: usize>(mut state: Self::State<D>, device: &B::Device) -> Self::State<D> {
    const LEARNING_RATE: LearningRate = 0.02;
    fn with_updated_params_should_have_state() {
    fn without_updated_params_should_not_have_state() {
    fn can_attach_gradient_clipping() {
    fn should_load_state() {
    fn random_tensor<B: Backend>(device: &B::Device) -> Tensor<B, 2> {
    fn layer<B: Backend>(device: &B::Device) -> Linear<B> {
    fn sgd_with_all()
    -> OptimizerAdaptor<Sgd<TestBackend>, Linear<TestAutodiffBackend>, TestAutodiffBackend> {
/// Wrapper struct that adapts any [simple optimizer](SimpleOptimizer) into
/// an [optimizer](Optimizer).
pub struct OptimizerAdaptor<O, M, B>
where
    O: SimpleOptimizer<B::InnerBackend>,
    M: AutodiffModule<B>,
    B: AutodiffBackend,
{
impl<O, B, M> From<O> for OptimizerAdaptor<O, M, B>
where
    B: AutodiffBackend,
    M: AutodiffModule<B>,
    O: SimpleOptimizer<B::InnerBackend>,
{
    fn from(optim: O) -> Self {
impl<O, M, B> OptimizerAdaptor<O, M, B>
where
    O: SimpleOptimizer<B::InnerBackend>,
    M: AutodiffModule<B>,
    B: AutodiffBackend,
{
/// Sets the gradient clipping.
///
/// # Arguments
///
/// * `gradient_clipping` - The gradient clipping.
///
/// # Returns
///
/// The optimizer.
    pub fn with_grad_clipping(mut self, gradient_clipping: GradientClipping) -> Self {
    pub(crate) fn has_gradient_clipping(&self) -> bool {
impl<O, B, M> Optimizer<M, B> for OptimizerAdaptor<O, M, B>
where
    B: AutodiffBackend,
    M: AutodiffModule<B>,
    O: SimpleOptimizer<B::InnerBackend>,
{
    type Record = HashMap<ParamId, AdaptorRecord<O, B>>;
    fn step(&mut self, lr: LearningRate, module: M, mut grads: GradientsParams) -> M {
    fn to_record(&self) -> Self::Record {
    fn load_record(mut self, record: Self::Record) -> Self {
struct SimpleOptimizerMapper<'a, M, B, O>
where
    M: AutodiffModule<B>,
    B: AutodiffBackend,
    O: SimpleOptimizer<B::InnerBackend>,
{
impl<M, B, O> ModuleMapper<B> for SimpleOptimizerMapper<'_, M, B, O>
where
    M: AutodiffModule<B>,
    B: AutodiffBackend,
    O: SimpleOptimizer<B::InnerBackend>,
{
    fn map_float<const D: usize>(&mut self, param: Param<Tensor<B, D>>) -> Param<Tensor<B, D>> {
/// Simple optimizer is an opinionated trait to simplify the process of implementing an
/// optimizer.
///
/// Implementations don't have to handle missing gradients, loading and exporting records, navigate the
/// module parameter structure, handle tracked and untracked tensors, and the likes.
pub trait SimpleOptimizer<B>: Send + Sync + Clone
where
    B: Backend,
{
/// The state of the optimizer. It also implements [record](Record), so that it can be saved.
    type State<const D: usize>: Record<B> + Clone + 'static;
/// The optimizer step is performed for one tensor at a time with its gradient and state.
///
/// Note that the state is passed as parameter, so implementations don't have to handle
/// the saving and loading of recorded states.
    fn step<const D: usize>(
        &self,
        lr: LearningRate,
        tensor: Tensor<B, D>,
        grad: Tensor<B, D>,
        state: Option<Self::State<D>>,
    ) -> (Tensor<B, D>, Option<Self::State<D>>);
/// Change the device of the state.
///
/// This function will be called accordingly to have the state on the same device as the
/// gradient and the tensor when the [step](SimpleOptimizer::step) function is called.
    fn to_device<const D: usize>(state: Self::State<D>, device: &B::Device) -> Self::State<D>;
/// Adaptor module for optimizers.
/// Record module for optimizers.
/// [Optimizer adaptor](crate::optim::simple::adaptor::OptimizerAdaptor) record.
///
/// Records are versioned for backward compatibility, so old records can be loaded.
pub enum AdaptorRecord<O, B>
where
    O: SimpleOptimizer<B::InnerBackend>,
    B: AutodiffBackend,
{
/// Version 1.
/// [Optimizer adaptor](crate::optim::simple::adaptor::OptimizerAdaptor) record item.
pub enum AdaptorRecordItem<
    O: SimpleOptimizer<B::InnerBackend>,
    B: AutodiffBackend,
    S: PrecisionSettings,
> {
/// Version 1.
impl<O, B> Record<B> for AdaptorRecord<O, B>
where
    O: SimpleOptimizer<B::InnerBackend>,
    B: AutodiffBackend,
{
    type Item<S: PrecisionSettings> = AdaptorRecordItem<O, B, S>;
    fn into_item<S: PrecisionSettings>(self) -> Self::Item<S> {
    fn from_item<S: PrecisionSettings>(item: Self::Item<S>, device: &B::Device) -> Self {
impl<O, B> Clone for AdaptorRecord<O, B>
where
    O: SimpleOptimizer<B::InnerBackend>,
    B: AutodiffBackend,
{
    fn clone(&self) -> Self {
impl<O, B> AdaptorRecord<O, B>
where
    O: SimpleOptimizer<B::InnerBackend>,
    B: AutodiffBackend,
{
/// Converts the record into the optimizer state.
///
/// # Returns
///
/// The optimizer state.
    pub fn into_state<const D: usize>(self) -> O::State<D> {
/// Converts the optimizer state into the record.
///
/// # Arguments
///
/// * `state`: The optimizer state.
///
/// # Returns
///
/// The record.
    pub fn from_state<const D: usize>(state: O::State<D>) -> Self {
/// [Optimizer adaptor](crate::optim::simple::adaptor::OptimizerAdaptor) record item.
pub enum AdaptorRecordV1<O: SimpleOptimizer<B>, B: Backend> {
/// Rank 0.
/// Rank 1.
/// Rank 2.
/// Rank 3.
/// Rank 4.
/// Rank 5.
/// Rank 6.
/// Rank 7.
/// Rank 8.
impl<O: SimpleOptimizer<B>, B: Backend> Clone for AdaptorRecordV1<O, B> {
    fn clone(&self) -> Self {
/// [Optimizer adaptor](crate::optim::simple::adaptor::OptimizerAdaptor) record item.
pub enum AdaptorRecordItemV1<O: SimpleOptimizer<B>, B: Backend, S: PrecisionSettings> {
/// Rank 0.
/// Rank 1.
/// Rank 2.
/// Rank 3.
/// Rank 4.
/// Rank 5.
/// Rank 6.
/// Rank 7.
/// Rank 8.
impl<O, B> AdaptorRecordV1<O, B>
where
    O: SimpleOptimizer<B>,
    B: Backend,
{
/// Convert the record into the state.
///
/// # Returns
///
/// The state.
///
/// # Panics
///
/// Panics if the state dimension is not supported.
    pub fn into_state<const D: usize>(self) -> O::State<D> {
/// Convert the state into the record.
///
/// # Arguments
///
/// * `state`: The state.
///
/// # Returns
///
/// The record.
    pub fn from_state<const D: usize>(state: O::State<D>) -> Self {
impl<O, B> Record<B> for AdaptorRecordV1<O, B>
where
    O: SimpleOptimizer<B>,
    B: Backend,
{
    type Item<S: PrecisionSettings> = AdaptorRecordItemV1<O, B, S>;
    fn into_item<S: PrecisionSettings>(self) -> Self::Item<S> {
    fn from_item<S: PrecisionSettings>(item: Self::Item<S>, device: &B::Device) -> Self {
pub struct GradientsParamsConverter<'a, M: AutodiffModule<B>, B: AutodiffBackend> {
pub struct GradientsParamsChangeDevice<'a, M: AutodiffModule<B>, B: AutodiffBackend> {
impl<B, M> ModuleVisitor<B> for GradientsParamsConverter<'_, M, B>
where
    B: AutodiffBackend,
    M: AutodiffModule<B>,
{
    fn visit_float<const D: usize>(&mut self, param: &Param<Tensor<B, D>>) {
impl<B, M> ModuleVisitor<B> for GradientsParamsChangeDevice<'_, M, B>
where
    B: AutodiffBackend,
    M: AutodiffModule<B>,
{
    fn visit_float<const D: usize>(&mut self, param: &Param<Tensor<B, D>>) {
pub struct WsClient {
impl WsClient {
    pub fn init(device: WsDevice) -> Self {
    pub(crate) fn new(
        device: WsDevice,
        sender: Sender<ClientRequest>,
        runtime: Arc<tokio::runtime::Runtime>,
        session_id: SessionId,
    ) -> Self {
pub(crate) struct WsSender {
impl WsSender {
    pub(crate) fn send(&self, task: ComputeTask) {
    pub(crate) fn new_tensor_id(&self) -> TensorId {
    pub(crate) fn send_callback(
        &self,
        task: ComputeTask,
    ) -> impl Future<Output = TaskResponseContent> + Send + use<> {
    pub(crate) fn close(&mut self) {
impl Drop for WsSender {
    fn drop(&mut self) {
/// A local channel with direct connection to the backend runner clients.
pub struct WsChannel;
impl RunnerChannel for WsChannel {
    type Device = WsDevice;
    type Bridge = WsBridge;
    type Client = WsClient;
    type FloatElem = f32;
    type IntElem = i32;
    type BoolElem = u32;
    fn name(device: &Self::Device) -> String {
    fn init_client(device: &Self::Device) -> Self::Client {
    fn get_tensor_handle(tensor: &TensorIr, client: &Self::Client) -> RemoteTensorHandle {
    fn register_tensor(
        client: &Self::Client,
        handle: RemoteTensorHandle,
        shape: Vec<usize>,
        dtype: burn_tensor::DType,
    ) -> RouterTensor<Self::Client> {
    fn change_client_backend(
        tensor: RouterTensor<Self::Client>,
        target_device: &Self::Device, // target device
    ) -> RouterTensor<Self::Client> {
// Get tensor handle from current client
// It is very important to block on any request made with the sender, since ordering is crucial
// when registering operation or creating tensors.
//
// The overhead is minimal, since we only wait for the task to be sent to the async
// channel, but not sent to the websocket server and even less processed by the server.
impl RunnerClient for WsClient {
    type Device = WsDevice;
    fn register(&self, op: burn_ir::OperationIr) {
    fn read_tensor(&self, tensor: burn_ir::TensorIr) -> DynFut<TensorData> {
// Important for ordering to call the creation of the future sync.
    fn register_tensor_data(&self, data: TensorData) -> RouterTensor<Self> {
    fn register_empty_tensor(
        &self,
        shape: Vec<usize>,
        dtype: burn_tensor::DType,
    ) -> RouterTensor<Self> {
    fn register_float_tensor(
        &self,
        shape: Vec<usize>,
        _dtype: burn_tensor::FloatDType,
    ) -> RouterTensor<Self> {
    fn device(&self) -> Self::Device {
    fn sync(&self) {
// Important for ordering to call the creation of the future sync.
    fn seed(&self, _seed: u64) {
// TODO
/// The device contains the connection information of the server.
pub struct WsDevice {
// Unique ID generated from hash of the address
impl WsDevice {
/// Create a device from an url.
    pub fn new(url: &str) -> Self {
impl Default for WsDevice {
    fn default() -> Self {
impl DeviceOps for WsDevice {
    fn id(&self) -> DeviceId {
pub struct WsBridge;
pub struct RemoteTensorHandle {
impl RemoteTensorHandle {
/// Changes the backend of the tensor via a WebSocket.
/// We ask the original server to expose the tensor, then ask the target server to fetch
/// the tensor. The target server will open a new websocket connection to the original server
/// to download the data.
/// This way the client never sees the tensor's data, and we avoid a bottleneck.
    pub(crate) fn change_backend(mut self, target_device: &WsDevice) -> Self {
impl MultiBackendBridge for WsBridge {
    type TensorHandle = RemoteTensorHandle;
    type Device = WsDevice;
    fn change_backend_float(
        tensor: Self::TensorHandle,
        _shape: burn_tensor::Shape,
        target_device: &Self::Device,
    ) -> Self::TensorHandle {
    fn change_backend_int(
        tensor: Self::TensorHandle,
        _shape: burn_tensor::Shape,
        target_device: &Self::Device,
    ) -> Self::TensorHandle {
    fn change_backend_bool(
        tensor: Self::TensorHandle,
        _shape: burn_tensor::Shape,
        target_device: &Self::Device,
    ) -> Self::TensorHandle {
pub type CallbackSender = async_channel::Sender<TaskResponseContent>;
pub enum ClientRequest {
pub(crate) struct ClientWorker {
impl ClientWorker {
    async fn on_response(&mut self, response: TaskResponse) {
    fn register_callback(&mut self, id: ConnectionId, callback: CallbackSender) {
impl ClientWorker {
    pub fn start(device: WsDevice) -> WsClient {
        const MB: usize = 1024 * 1024;
// 64 KiB (previous default)
// 64 KiB (previous default)
// Init the connection.
// Websocket async worker loading callback from the server.
// Channel async worker sending operations to the server.
/// The remote backend allows you to run computation on a remote device.
///
/// Make sure there is a running server before trying to connect to it.
///
/// ```rust, ignore
/// fn main() {
///     let device = Default::default();
///     let port = 3000;
///
///     // You need to activate the `server` feature flag to have access to this function.
///     burn::server::start::<burn::backend::Wgpu>(device, port);
/// }
///```
    pub type RemoteBackend = BackendRouter<WsChannel>;
    pub fn test_to_device_over_websocket() {
// Some random input
// Move tensor to device 2
// Move tensor back to device 1
pub struct WsServer<B: BackendIr> {
impl<B: BackendIr> WsServer<B> {
/// Start the server on the given address.
    pub async fn start(device: Device<B>, port: u16) {
// The wgpu crate is logging too much, so we skip `info` level.
// If we start multiple servers in the same process, this will fail, it's ok
// build our application with some routes
// run it with hyper
    async fn handler_response(
        ws: WebSocketUpgrade,
        State(state): State<Self>,
    ) -> impl IntoResponse {
    async fn handler_request(ws: WebSocketUpgrade, State(state): State<Self>) -> impl IntoResponse {
    async fn handler_data(ws: WebSocketUpgrade, State(state): State<Self>) -> impl IntoResponse {
    async fn handle_socket_response(self, mut socket: WebSocket) {
    async fn handle_socket_request(self, mut socket: WebSocket) {
    async fn handle_socket_data(self, mut socket: WebSocket) {
// Get the requested exposed tensor data
// Send tensor and increment its counter
    async fn get_exposed_tensor_bytes(id: TensorId, state: Arc<TensorDataService>) -> bytes::Bytes {
// take the tensor out of the hashmap while we download
//panic!("A tensor was requested (id: {id:?}) that isn't being served");
// TODO should i add a
// Optional: Add a small pause to reduce CPU usage
// thread::yield_now();
pub async fn start_async<B: BackendIr>(device: Device<B>, port: u16) {
/// Start the server on the given port and [device](Device).
pub async fn start<B: BackendIr>(device: Device<B>, port: u16) {
/// The goal of the processor is to asynchronously process compute tasks on it own thread.
pub struct Processor<B: BackendIr> {
pub type Callback<M> = Sender<M>;
pub enum ProcessorTask {
impl<B: BackendIr> Processor<B> {
    pub async fn start(runner: Runner<B>, state: Arc<TensorDataService>) -> Sender<ProcessorTask> {
// channel for tasks to execute
/// Downloads a tensor that is exposed on another server. Requires a Tokio 1.x runtime
    async fn download_tensor(remote_tensor: TensorRemote) -> Option<TensorData> {
        const MB: usize = 1024 * 1024;
// 64 KiB (previous default)
/// A session manager control the creation of sessions.
///
/// Each session manages its own stream, spawning one thread per stream to mimic the same behavior
/// a native backend would have.
pub struct SessionManager<B: BackendIr> {
struct Session<B: BackendIr> {
impl<B: BackendIr> SessionManager<B> {
    pub fn new(device: Device<B>, server_state: Arc<TensorDataService>) -> Self {
/// Register a new responder for the session. Only one responder can exist for a session for
/// now.
    pub async fn register_responder(
        &self,
        session_id: SessionId,
    ) -> Receiver<Receiver<TaskResponse>> {
/// Get the stream for the current session and task.
    pub async fn stream(
        &self,
        session_id: &mut Option<SessionId>,
        task: Task,
    ) -> Option<(Stream<B>, ConnectionId, ComputeTask)> {
/// Close the session with the given id.
    pub async fn close(&self, session_id: Option<SessionId>) {
    fn register_session(&self, sessions: &mut HashMap<SessionId, Session<B>>, id: SessionId) {
impl<B: BackendIr> Session<B> {
    fn new(runner: Runner<B>, server_state: Arc<TensorDataService>) -> Self {
    fn init_responder(&mut self) -> Receiver<Receiver<TaskResponse>> {
/// Select the current [stream](Stream) based on the given task.
    async fn select(&mut self, stream_id: StreamId) -> Stream<B> {
// We return the stream.
// Close all streams created in the session.
    async fn close(&mut self) {
/// A stream makes sure all operations registered are executed in the order they were sent to the
/// server, protentially waiting to reconstruct consistency.
pub struct Stream<B: BackendIr> {
impl<B: BackendIr> Stream<B> {
    pub async fn new(
        runner: Runner<B>,
        writer_sender: Sender<Receiver<TaskResponse>>,
        state: Arc<TensorDataService>,
    ) -> Self {
    pub async fn register_operation(&self, op: Box<OperationIr>) {
    pub async fn register_tensor(&self, tensor_id: TensorId, data: TensorData) {
    pub async fn register_tensor_remote(&self, tensor: TensorRemote, new_id: TensorId) {
    pub async fn expose_tensor_remote(&self, tensor: TensorIr, count: u32) {
    pub async fn read_tensor(&self, id: ConnectionId, desc: TensorIr) {
    pub async fn sync(&self, id: ConnectionId) {
    pub async fn close(&self) {
pub struct TensorExposeState {
pub struct TensorDataService {
pub struct ConnectionId {
/// Unique identifier that can represent a session.
pub struct SessionId {
impl Display for SessionId {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
impl SessionId {
/// Create a new [session id](SessionId).
    pub fn new() -> Self {
pub enum Task {
pub struct TensorRemote {
pub enum ComputeTask {
/// Used by a server to request a tensor from another server
pub struct RemoteTensorReq {
pub struct TaskResponse {
pub enum TaskResponseContent {
pub type Rocm<F = f32, I = i32, B = u8> = CubeBackend<HipRuntime, F, I, B>;
pub type Rocm<F = f32, I = i32, B = u8> = burn_fusion::Fusion<CubeBackend<HipRuntime, F, I, B>>;
    pub type TestRuntime = cubecl::hip::HipRuntime;
// TODO: Add tests for bf16
// burn_cubecl::testgen_all!([f16, f32], [i8, i16, i32, i64], [u8, u32]);
/// A backend that forwards the tensor operations to the appropriate backend (given multiple backends).
pub struct BackendRouter<R: RunnerChannel> {
impl<R: RunnerChannel> core::fmt::Debug for BackendRouter<R> {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
impl<R: RunnerChannel> Clone for BackendRouter<R> {
    fn clone(&self) -> Self {
impl<R: RunnerChannel> Default for BackendRouter<R> {
    fn default() -> Self {
// TODO: quantization tensor primitive (w/ qparams)
impl<R: RunnerClient> QTensorPrimitive for RouterTensor<R> {
    fn scheme(&self) -> &QuantScheme {
impl<R: RunnerChannel> Backend for BackendRouter<R> {
    type Device = R::Device;
    type FloatTensorPrimitive = RouterTensor<R::Client>;
    type FloatElem = R::FloatElem;
    type IntTensorPrimitive = RouterTensor<R::Client>;
    type IntElem = R::IntElem;
    type BoolTensorPrimitive = RouterTensor<R::Client>;
    type BoolElem = R::BoolElem;
    type QuantizedTensorPrimitive = RouterTensor<R::Client>;
    type QuantizedEncoding = u32;
    fn name(device: &Self::Device) -> String {
    fn seed(seed: u64) {
    fn sync(device: &Self::Device) {
/// Allows tensors to be transferred between multiple backends.
pub trait MultiBackendBridge: Send + Sync + 'static {
/// The type that can be used to point to a tensor of any kind.
    type TensorHandle;
/// Device type used by the backends.
    type Device: DeviceOps;
/// Change the backend of the given float tensor.
    fn change_backend_float(
        tensor: Self::TensorHandle,
        shape: Shape,
        target_device: &Self::Device,
    ) -> Self::TensorHandle;
/// Change the backend of the given int tensor.
    fn change_backend_int(
        tensor: Self::TensorHandle,
        shape: Shape,
        target_device: &Self::Device,
    ) -> Self::TensorHandle;
/// Change the backend of the given bool tensor.
    fn change_backend_bool(
        tensor: Self::TensorHandle,
        shape: Shape,
        target_device: &Self::Device,
    ) -> Self::TensorHandle;
// TODO: change_backend_quantized
/// Simply transfers tensors between backends via the underlying [tensor data](burn_tensor::TensorData).
pub struct ByteBridge<Backends> {
/// Type alias for `<Br as MultiBackendBridge>::TensorHandle`.
pub type TensorHandle<Br> = <Br as MultiBackendBridge>::TensorHandle;
/// Defines the connection channel and operations for a setup with multiple backend runner clients.
pub trait RunnerChannel: Clone + Send + Sync + 'static + Sized {
/// Device type.
    type Device: DeviceOps;
/// A bridge that can transfer tensors between multiple backends.
    type Bridge: MultiBackendBridge<Device = Self::Device>;
/// Client type.
    type Client: RunnerClient<Device = Self::Device>;
/// Float element type.
    type FloatElem: Element;
/// Int element type.
    type IntElem: Element;
/// Bool element type.
    type BoolElem: Element;
/// Name of the channel.
    fn name(device: &Self::Device) -> String;
/// Initialize a new client for the given device.
    fn init_client(device: &Self::Device) -> Self::Client;
/// Get the tensor handle corresponding to the [tensor representation](TensorIr).
    fn get_tensor_handle(tensor: &TensorIr, client: &Self::Client) -> TensorHandle<Self::Bridge>;
/// Create a tensor with the given handle and shape.
    fn register_tensor(
        client: &Self::Client,
        handle: TensorHandle<Self::Bridge>,
        shape: Vec<usize>,
        dtype: DType,
    ) -> RouterTensor<Self::Client>;
/// Change the tensor to a different client backend.
    fn change_client_backend(
        tensor: RouterTensor<Self::Client>,
        device: &Self::Device, // target device
    ) -> RouterTensor<Self::Client> {
// Get tensor handle from current client
// Register tensor handle on target client
/// A local channel with direct connection to the backend runner clients.
pub struct DirectChannel<Backends, Bridge> {
impl<Backends, Bridge> Clone for DirectChannel<Backends, Bridge> {
    fn clone(&self) -> Self {
/// Type alias for `<R as RunnerChannel>::Client`.
pub type Client<R> = <R as RunnerChannel>::Client;
pub(crate) static CLIENTS: RunnerClientLocator = RunnerClientLocator::new();
static SEED_SET: AtomicBool = AtomicBool::new(false);
static SEED: AtomicU64 = AtomicU64::new(0);
type Key = (core::any::TypeId, DeviceId);
/// Define how to interact with the runner.
pub trait RunnerClient: Clone + Send + Sync + Sized {
/// Device type.
    type Device: DeviceOps;
/// Register a new tensor operation to be executed by the (runner) server.
    fn register(&self, op: OperationIr);
/// Read the values contained by a tensor.
    fn read_tensor(&self, tensor: TensorIr) -> DynFut<TensorData>;
/// Sync the runner, ensure that all computations are finished.
    fn sync(&self);
/// Create a new [RouterTensor] from the tensor data.
    fn register_tensor_data(&self, data: TensorData) -> RouterTensor<Self>;
/// Create a new [RouterTensor] with no resources associated.
    fn register_empty_tensor(&self, shape: Vec<usize>, dtype: DType) -> RouterTensor<Self>;
/// Create a new float [RouterTensor] with no resources associated.
    fn register_float_tensor(&self, shape: Vec<usize>, dtype: FloatDType) -> RouterTensor<Self>;
/// Get the current device used by all operations handled by this client.
    fn device(&self) -> Self::Device;
/// Seed the runner.
    fn seed(&self, seed: u64);
pub(crate) struct RunnerClientLocator {
/// Get the client for the given device
pub fn get_client<R: RunnerChannel>(device: &R::Device) -> Client<R> {
pub(crate) fn set_seed(seed: u64) {
fn get_seed() -> Option<u64> {
/// Initialize a new client for the given device.
///
/// If a (global) seed was previously set, the client seed is set.
fn new_client<R: RunnerChannel>(device: &R::Device) -> Client<R> {
impl RunnerClientLocator {
/// Create a new client locator.
    pub const fn new() -> Self {
/// Get the runner client for the given device.
///
/// If a client isn't already initialized, it is created.
    pub fn client<R: RunnerChannel + 'static>(&self, device: &R::Device) -> Client<R> {
    fn register_inner<R: RunnerChannel + 'static>(
        key: Key,
        client: Client<R>,
        clients: &mut Option<HashMap<Key, Box<dyn core::any::Any + Send>>>,
    ) {
//! Burn multi-backend router.
/// A local channel with a simple byte bridge between backends.
/// It transfers tensors between backends via the underlying [tensor data](burn_tensor::TensorData).
pub type DirectByteChannel<Backends> = DirectChannel<Backends, ByteBridge<Backends>>;
/// Router backend.
///
/// # Example
///
/// ```ignore
/// type MyBackend = Router<(NdArray, Wgpu)>;
/// ```
pub type Router<Backends> = BackendRouter<DirectByteChannel<Backends>>;
    pub type TestBackend1 = burn_ndarray::NdArray<f32, i32>;
    pub type TestBackend2 = burn_wgpu::Wgpu<f32, i32>;
    pub type TestBackend = BackendRouter<DirectByteChannel<(TestBackend1, TestBackend2)>>;
    pub type TestTensor<const D: usize> = burn_tensor::Tensor<TestBackend, D>;
    pub type TestTensorInt<const D: usize> = burn_tensor::Tensor<TestBackend, D, burn_tensor::Int>;
    pub type TestTensorBool<const D: usize> =
        burn_tensor::Tensor<TestBackend, D, burn_tensor::Bool>;
// TODO: add support for quantization
// burn_tensor::testgen_quantization!();
impl<R: RunnerChannel> ActivationOps<Self> for BackendRouter<R> {
impl<R: RunnerChannel> BoolTensorOps<Self> for BackendRouter<R> {
    fn bool_empty(shape: Shape, device: &Device<Self>) -> BoolTensor<Self> {
// Get the runtime client on which to register the operation for execution.
    async fn bool_into_data(tensor: BoolTensor<Self>) -> TensorData {
    fn bool_from_data(data: TensorData, device: &Device<Self>) -> BoolTensor<Self> {
    fn bool_into_int(tensor: BoolTensor<Self>) -> IntTensor<Self> {
    fn bool_into_float(tensor: BoolTensor<Self>) -> FloatTensor<Self> {
    fn bool_device(tensor: &BoolTensor<Self>) -> Device<Self> {
    fn bool_to_device(tensor: BoolTensor<Self>, device: &Device<Self>) -> BoolTensor<Self> {
    fn bool_reshape(tensor: BoolTensor<Self>, shape: Shape) -> BoolTensor<Self> {
    fn bool_slice(
        tensor: BoolTensor<Self>,
        ranges: &[core::ops::Range<usize>],
    ) -> BoolTensor<Self> {
    fn bool_slice_assign(
        tensor: BoolTensor<Self>,
        ranges: &[core::ops::Range<usize>],
        value: BoolTensor<Self>,
    ) -> BoolTensor<Self> {
    fn bool_equal(lhs: BoolTensor<Self>, rhs: BoolTensor<Self>) -> BoolTensor<Self> {
    fn bool_not(tensor: BoolTensor<Self>) -> BoolTensor<Self> {
    fn bool_and(lhs: BoolTensor<Self>, rhs: BoolTensor<Self>) -> BoolTensor<Self> {
    fn bool_or(lhs: BoolTensor<Self>, rhs: BoolTensor<Self>) -> BoolTensor<Self> {
    fn bool_swap_dims(tensor: BoolTensor<Self>, dim1: usize, dim2: usize) -> BoolTensor<Self> {
    fn bool_permute(tensor: BoolTensor<Self>, axes: &[usize]) -> BoolTensor<Self> {
// Change the shape of the tensor to match the new axes
    fn bool_flip(tensor: BoolTensor<Self>, axes: &[usize]) -> BoolTensor<Self> {
    fn bool_expand(tensor: BoolTensor<Self>, shape: Shape) -> BoolTensor<Self> {
    fn bool_cat(tensors: Vec<BoolTensor<Self>>, dim: usize) -> BoolTensor<Self> {
// Calculate the output shape
    fn bool_repeat_dim(tensor: BoolTensor<Self>, dim: usize, times: usize) -> BoolTensor<Self> {
impl<R: RunnerChannel> FloatTensorOps<Self> for BackendRouter<R> {
    fn float_from_data(data: TensorData, device: &Device<Self>) -> FloatTensor<Self> {
    fn float_random(
        shape: Shape,
        distribution: Distribution,
        device: &Device<Self>,
    ) -> FloatTensor<Self> {
// Get the runtime client on which to register the operation for execution.
    fn float_zeros(shape: Shape, device: &Device<Self>) -> FloatTensor<Self> {
// Get the runtime client on which to register the operation for execution.
    fn float_ones(shape: Shape, device: &Device<Self>) -> FloatTensor<Self> {
// Get the runtime client on which to register the operation for execution.
    fn float_full(
        shape: Shape,
        fill_value: FloatElem<Self>,
        device: &Device<Self>,
    ) -> FloatTensor<Self> {
// Get the runtime client on which to register the operation for execution.
    async fn float_into_data(tensor: FloatTensor<Self>) -> TensorData {
// Since underlying backends can have different data types, we convert to the current elem
    fn float_device(tensor: &FloatTensor<Self>) -> Device<Self> {
    fn float_to_device(tensor: FloatTensor<Self>, device: &Device<Self>) -> FloatTensor<Self> {
    fn float_into_int(tensor: FloatTensor<Self>) -> IntTensor<Self> {
    fn float_empty(shape: Shape, device: &Device<Self>) -> FloatTensor<Self> {
// Get the runtime client on which to register the operation for execution.
    fn float_add(lhs: FloatTensor<Self>, rhs: FloatTensor<Self>) -> FloatTensor<Self> {
    fn float_add_scalar(lhs: FloatTensor<Self>, rhs: FloatElem<Self>) -> FloatTensor<Self> {
    fn float_clamp(
        tensor: FloatTensor<Self>,
        min: FloatElem<Self>,
        max: FloatElem<Self>,
    ) -> FloatTensor<Self> {
    fn float_sub(lhs: FloatTensor<Self>, rhs: FloatTensor<Self>) -> FloatTensor<Self> {
    fn float_sub_scalar(lhs: FloatTensor<Self>, rhs: FloatElem<Self>) -> FloatTensor<Self> {
    fn float_mul(lhs: FloatTensor<Self>, rhs: FloatTensor<Self>) -> FloatTensor<Self> {
    fn float_mul_scalar(lhs: FloatTensor<Self>, rhs: FloatElem<Self>) -> FloatTensor<Self> {
    fn float_div(lhs: FloatTensor<Self>, rhs: FloatTensor<Self>) -> FloatTensor<Self> {
    fn float_div_scalar(lhs: FloatTensor<Self>, rhs: FloatElem<Self>) -> FloatTensor<Self> {
    fn float_remainder(lhs: FloatTensor<Self>, rhs: FloatTensor<Self>) -> FloatTensor<Self> {
    fn float_remainder_scalar(lhs: FloatTensor<Self>, rhs: FloatElem<Self>) -> FloatTensor<Self> {
    fn float_matmul(lhs: FloatTensor<Self>, rhs: FloatTensor<Self>) -> FloatTensor<Self> {
    fn float_swap_dims(tensor: FloatTensor<Self>, dim1: usize, dim2: usize) -> FloatTensor<Self> {
    fn float_reshape(tensor: FloatTensor<Self>, shape: Shape) -> FloatTensor<Self> {
    fn float_gather(
        dim: usize,
        tensor: FloatTensor<Self>,
        indices: IntTensor<Self>,
    ) -> FloatTensor<Self> {
    fn float_scatter(
        dim: usize,
        tensor: FloatTensor<Self>,
        indices: IntTensor<Self>,
        value: FloatTensor<Self>,
    ) -> FloatTensor<Self> {
    fn float_select(
        tensor: FloatTensor<Self>,
        dim: usize,
        indices: IntTensor<Self>,
    ) -> FloatTensor<Self> {
    fn float_select_assign(
        tensor: FloatTensor<Self>,
        dim: usize,
        indices: IntTensor<Self>,
        value: FloatTensor<Self>,
    ) -> FloatTensor<Self> {
    fn float_slice(tensor: FloatTensor<Self>, ranges: &[Range<usize>]) -> FloatTensor<Self> {
    fn float_slice_assign(
        tensor: FloatTensor<Self>,
        ranges: &[Range<usize>],
        value: FloatTensor<Self>,
    ) -> FloatTensor<Self> {
    fn float_mask_where(
        tensor: FloatTensor<Self>,
        mask: BoolTensor<Self>,
        value: FloatTensor<Self>,
    ) -> FloatTensor<Self> {
    fn float_mask_fill(
        tensor: FloatTensor<Self>,
        mask: BoolTensor<Self>,
        value: FloatElem<Self>,
    ) -> FloatTensor<Self> {
    fn float_equal(lhs: FloatTensor<Self>, rhs: FloatTensor<Self>) -> BoolTensor<Self> {
    fn float_equal_elem(lhs: FloatTensor<Self>, rhs: FloatElem<Self>) -> BoolTensor<Self> {
    fn float_greater(lhs: FloatTensor<Self>, rhs: FloatTensor<Self>) -> BoolTensor<Self> {
    fn float_greater_elem(lhs: FloatTensor<Self>, rhs: FloatElem<Self>) -> BoolTensor<Self> {
    fn float_greater_equal(lhs: FloatTensor<Self>, rhs: FloatTensor<Self>) -> BoolTensor<Self> {
    fn float_greater_equal_elem(lhs: FloatTensor<Self>, rhs: FloatElem<Self>) -> BoolTensor<Self> {
    fn float_lower(lhs: FloatTensor<Self>, rhs: FloatTensor<Self>) -> BoolTensor<Self> {
    fn float_lower_elem(lhs: FloatTensor<Self>, rhs: FloatElem<Self>) -> BoolTensor<Self> {
    fn float_lower_equal(lhs: FloatTensor<Self>, rhs: FloatTensor<Self>) -> BoolTensor<Self> {
    fn float_lower_equal_elem(lhs: FloatTensor<Self>, rhs: FloatElem<Self>) -> BoolTensor<Self> {
    fn float_sum(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
    fn float_sum_dim(tensor: FloatTensor<Self>, axis: usize) -> FloatTensor<Self> {
    fn float_prod(tensor: IntTensor<Self>) -> IntTensor<Self> {
    fn float_prod_dim(tensor: IntTensor<Self>, dim: usize) -> IntTensor<Self> {
    fn float_mean(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
    fn float_mean_dim(tensor: FloatTensor<Self>, dim: usize) -> FloatTensor<Self> {
    fn float_exp(lhs: FloatTensor<Self>) -> FloatTensor<Self> {
    fn float_log(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
    fn float_log1p(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
    fn float_powf_scalar(lhs: FloatTensor<Self>, rhs: f32) -> FloatTensor<Self> {
    fn float_sqrt(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
    fn float_abs(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
    fn float_cos(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
    fn float_sin(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
    fn float_tanh(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
    fn float_round(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
    fn float_floor(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
    fn float_ceil(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
    fn float_recip(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
    fn float_erf(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
    fn float_cat(tensors: Vec<FloatTensor<Self>>, dim: usize) -> FloatTensor<Self> {
// Calculate the output shape
    fn float_argmax(tensor: FloatTensor<Self>, dim: usize) -> IntTensor<Self> {
    fn float_repeat_dim(tensor: FloatTensor<Self>, dim: usize, times: usize) -> FloatTensor<Self> {
    fn float_argmin(tensor: FloatTensor<Self>, dim: usize) -> IntTensor<Self> {
    fn float_max(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
    fn float_max_dim(tensor: FloatTensor<Self>, dim: usize) -> FloatTensor<Self> {
    fn float_max_dim_with_indices(
        tensor: FloatTensor<Self>,
        dim: usize,
    ) -> (FloatTensor<Self>, IntTensor<Self>) {
    fn float_min(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
    fn float_min_dim(tensor: FloatTensor<Self>, dim: usize) -> FloatTensor<Self> {
    fn float_min_dim_with_indices(
        tensor: FloatTensor<Self>,
        dim: usize,
    ) -> (FloatTensor<Self>, IntTensor<Self>) {
    fn float_powf(lhs: FloatTensor<Self>, rhs: FloatTensor<Self>) -> FloatTensor<Self> {
    fn float_permute(tensor: FloatTensor<Self>, axes: &[usize]) -> FloatTensor<Self> {
// Change the shape of the tensor to match the new axes
    fn float_expand(tensor: FloatTensor<Self>, shape: Shape) -> FloatTensor<Self> {
    fn float_flip(tensor: FloatTensor<Self>, axes: &[usize]) -> FloatTensor<Self> {
    fn float_cast(tensor: FloatTensor<Self>, dtype: burn_tensor::FloatDType) -> FloatTensor<Self> {
impl<R: RunnerChannel> IntTensorOps<Self> for BackendRouter<R> {
    fn int_empty(shape: Shape, device: &Device<Self>) -> IntTensor<Self> {
// Get the runtime client on which to register the operation for execution.
    async fn int_into_data(tensor: IntTensor<Self>) -> TensorData {
// Since underlying backends can have different data types, we convert to the current elem
    fn int_from_data(data: TensorData, device: &Device<Self>) -> IntTensor<Self> {
    fn int_device(tensor: &IntTensor<Self>) -> Device<Self> {
    fn int_to_device(tensor: IntTensor<Self>, device: &Device<Self>) -> IntTensor<Self> {
    fn int_reshape(tensor: IntTensor<Self>, shape: Shape) -> IntTensor<Self> {
    fn int_slice(tensor: IntTensor<Self>, ranges: &[Range<usize>]) -> IntTensor<Self> {
    fn int_slice_assign(
        tensor: IntTensor<Self>,
        ranges: &[Range<usize>],
        value: IntTensor<Self>,
    ) -> IntTensor<Self> {
    fn int_mask_where(
        tensor: IntTensor<Self>,
        mask: BoolTensor<Self>,
        value: IntTensor<Self>,
    ) -> IntTensor<Self> {
    fn int_mask_fill(
        tensor: IntTensor<Self>,
        mask: BoolTensor<Self>,
        value: IntElem<Self>,
    ) -> IntTensor<Self> {
    fn int_gather(
        dim: usize,
        tensor: IntTensor<Self>,
        indices: IntTensor<Self>,
    ) -> IntTensor<Self> {
    fn int_scatter(
        dim: usize,
        tensor: IntTensor<Self>,
        indices: IntTensor<Self>,
        value: IntTensor<Self>,
    ) -> IntTensor<Self> {
    fn int_select(
        tensor: IntTensor<Self>,
        dim: usize,
        indices: IntTensor<Self>,
    ) -> IntTensor<Self> {
    fn int_select_assign(
        tensor: IntTensor<Self>,
        dim: usize,
        indices: IntTensor<Self>,
        value: IntTensor<Self>,
    ) -> IntTensor<Self> {
    fn int_cat(tensors: Vec<IntTensor<Self>>, dim: usize) -> IntTensor<Self> {
// Calculate the output shape
    fn int_equal(lhs: IntTensor<Self>, rhs: IntTensor<Self>) -> BoolTensor<Self> {
    fn int_equal_elem(lhs: IntTensor<Self>, rhs: IntElem<Self>) -> BoolTensor<Self> {
    fn int_greater(lhs: IntTensor<Self>, rhs: IntTensor<Self>) -> BoolTensor<Self> {
    fn int_greater_elem(lhs: IntTensor<Self>, rhs: IntElem<Self>) -> BoolTensor<Self> {
    fn int_greater_equal(lhs: IntTensor<Self>, rhs: IntTensor<Self>) -> BoolTensor<Self> {
    fn int_greater_equal_elem(lhs: IntTensor<Self>, rhs: IntElem<Self>) -> BoolTensor<Self> {
    fn int_lower(lhs: IntTensor<Self>, rhs: IntTensor<Self>) -> BoolTensor<Self> {
    fn int_lower_elem(lhs: IntTensor<Self>, rhs: IntElem<Self>) -> BoolTensor<Self> {
    fn int_lower_equal(lhs: IntTensor<Self>, rhs: IntTensor<Self>) -> BoolTensor<Self> {
    fn int_lower_equal_elem(lhs: IntTensor<Self>, rhs: IntElem<Self>) -> BoolTensor<Self> {
    fn int_add(lhs: IntTensor<Self>, rhs: IntTensor<Self>) -> IntTensor<Self> {
    fn int_add_scalar(lhs: IntTensor<Self>, rhs: IntElem<Self>) -> IntTensor<Self> {
    fn int_sub(lhs: IntTensor<Self>, rhs: IntTensor<Self>) -> IntTensor<Self> {
    fn int_sub_scalar(lhs: IntTensor<Self>, rhs: IntElem<Self>) -> IntTensor<Self> {
    fn int_mul(lhs: IntTensor<Self>, rhs: IntTensor<Self>) -> IntTensor<Self> {
    fn int_mul_scalar(lhs: IntTensor<Self>, rhs: IntElem<Self>) -> IntTensor<Self> {
    fn int_div(lhs: IntTensor<Self>, rhs: IntTensor<Self>) -> IntTensor<Self> {
    fn int_div_scalar(lhs: IntTensor<Self>, rhs: IntElem<Self>) -> IntTensor<Self> {
    fn int_remainder(lhs: IntTensor<Self>, rhs: IntTensor<Self>) -> IntTensor<Self> {
    fn int_remainder_scalar(lhs: IntTensor<Self>, rhs: IntElem<Self>) -> IntTensor<Self> {
    fn int_zeros(shape: Shape, device: &Device<Self>) -> IntTensor<Self> {
// Get the runtime client on which to register the operation for execution.
    fn int_ones(shape: Shape, device: &Device<Self>) -> IntTensor<Self> {
// Get the runtime client on which to register the operation for execution.
    fn int_sum(tensor: IntTensor<Self>) -> IntTensor<Self> {
    fn int_sum_dim(tensor: IntTensor<Self>, axis: usize) -> IntTensor<Self> {
    fn int_prod(tensor: IntTensor<Self>) -> IntTensor<Self> {
    fn int_prod_dim(tensor: IntTensor<Self>, dim: usize) -> IntTensor<Self> {
    fn int_mean(tensor: IntTensor<Self>) -> IntTensor<Self> {
    fn int_mean_dim(tensor: IntTensor<Self>, dim: usize) -> IntTensor<Self> {
    fn int_argmax(tensor: IntTensor<Self>, dim: usize) -> IntTensor<Self> {
    fn int_argmin(tensor: IntTensor<Self>, dim: usize) -> IntTensor<Self> {
    fn int_clamp(
        tensor: IntTensor<Self>,
        min: IntElem<Self>,
        max: IntElem<Self>,
    ) -> IntTensor<Self> {
    fn int_abs(tensor: IntTensor<Self>) -> IntTensor<Self> {
    fn int_into_float(tensor: IntTensor<Self>) -> FloatTensor<Self> {
    fn int_swap_dims(tensor: IntTensor<Self>, dim1: usize, dim2: usize) -> IntTensor<Self> {
    fn int_max(tensor: IntTensor<Self>) -> IntTensor<Self> {
    fn int_max_dim(tensor: IntTensor<Self>, dim: usize) -> IntTensor<Self> {
    fn int_max_dim_with_indices(
        tensor: IntTensor<Self>,
        dim: usize,
    ) -> (IntTensor<Self>, IntTensor<Self>) {
    fn int_max_abs(tensor: IntTensor<Self>) -> IntTensor<Self> {
    fn int_max_abs_dim(tensor: IntTensor<Self>, dim: usize) -> IntTensor<Self> {
    fn int_min(tensor: IntTensor<Self>) -> IntTensor<Self> {
    fn int_min_dim(tensor: IntTensor<Self>, dim: usize) -> IntTensor<Self> {
    fn int_min_dim_with_indices(
        tensor: IntTensor<Self>,
        dim: usize,
    ) -> (IntTensor<Self>, IntTensor<Self>) {
    fn int_random(
        shape: Shape,
        distribution: Distribution,
        device: &Device<Self>,
    ) -> IntTensor<Self> {
// Get the runtime client on which to register the operation for execution.
    fn int_permute(tensor: IntTensor<Self>, axes: &[usize]) -> IntTensor<Self> {
// Change the shape of the tensor to match the new axes
    fn int_expand(tensor: IntTensor<Self>, shape: Shape) -> IntTensor<Self> {
    fn int_flip(tensor: IntTensor<Self>, axes: &[usize]) -> IntTensor<Self> {
    fn int_repeat_dim(tensor: IntTensor<Self>, dim: usize, times: usize) -> IntTensor<Self> {
    fn bitwise_and(lhs: IntTensor<Self>, rhs: IntTensor<Self>) -> IntTensor<Self> {
    fn bitwise_or(lhs: IntTensor<Self>, rhs: IntTensor<Self>) -> IntTensor<Self> {
    fn bitwise_xor(lhs: IntTensor<Self>, rhs: IntTensor<Self>) -> IntTensor<Self> {
    fn bitwise_not(tensor: IntTensor<Self>) -> IntTensor<Self> {
    fn bitwise_and_scalar(lhs: IntTensor<Self>, rhs: IntElem<Self>) -> IntTensor<Self> {
    fn bitwise_or_scalar(lhs: IntTensor<Self>, rhs: IntElem<Self>) -> IntTensor<Self> {
    fn bitwise_xor_scalar(lhs: IntTensor<Self>, rhs: IntElem<Self>) -> IntTensor<Self> {
    fn bitwise_left_shift(lhs: IntTensor<Self>, rhs: IntTensor<Self>) -> IntTensor<Self> {
    fn bitwise_left_shift_scalar(lhs: IntTensor<Self>, rhs: IntElem<Self>) -> IntTensor<Self> {
    fn bitwise_right_shift(lhs: IntTensor<Self>, rhs: IntTensor<Self>) -> IntTensor<Self> {
    fn bitwise_right_shift_scalar(lhs: IntTensor<Self>, rhs: IntElem<Self>) -> IntTensor<Self> {
impl<R: RunnerChannel> ModuleOps<Self> for BackendRouter<R> {
    fn conv1d(
        x: FloatTensor<Self>,
        weight: FloatTensor<Self>,
        bias: Option<FloatTensor<Self>>,
        options: ConvOptions<1>,
    ) -> FloatTensor<Self> {
    fn conv2d(
        x: FloatTensor<Self>,
        weight: FloatTensor<Self>,
        bias: Option<FloatTensor<Self>>,
        options: ConvOptions<2>,
    ) -> FloatTensor<Self> {
    fn conv3d(
        x: FloatTensor<Self>,
        weight: FloatTensor<Self>,
        bias: Option<FloatTensor<Self>>,
        options: ConvOptions<3>,
    ) -> FloatTensor<Self> {
    fn conv_transpose1d(
        x: FloatTensor<Self>,
        weight: FloatTensor<Self>,
        bias: Option<FloatTensor<Self>>,
        options: ConvTransposeOptions<1>,
    ) -> FloatTensor<Self> {
    fn conv_transpose2d(
        x: FloatTensor<Self>,
        weight: FloatTensor<Self>,
        bias: Option<FloatTensor<Self>>,
        options: ConvTransposeOptions<2>,
    ) -> FloatTensor<Self> {
    fn conv_transpose3d(
        x: FloatTensor<Self>,
        weight: FloatTensor<Self>,
        bias: Option<FloatTensor<Self>>,
        options: ConvTransposeOptions<3>,
    ) -> FloatTensor<Self> {
    fn avg_pool1d(
        x: FloatTensor<Self>,
        kernel_size: usize,
        stride: usize,
        padding: usize,
        count_include_pad: bool,
    ) -> FloatTensor<Self> {
    fn avg_pool2d(
        x: FloatTensor<Self>,
        kernel_size: [usize;
    fn avg_pool1d_backward(
        x: FloatTensor<Self>,
        grad: FloatTensor<Self>,
        kernel_size: usize,
        stride: usize,
        padding: usize,
        count_include_pad: bool,
    ) -> FloatTensor<Self> {
    fn avg_pool2d_backward(
        x: FloatTensor<Self>,
        grad: FloatTensor<Self>,
        kernel_size: [usize;
    fn max_pool1d(
        x: FloatTensor<Self>,
        kernel_size: usize,
        stride: usize,
        padding: usize,
        dilation: usize,
    ) -> FloatTensor<Self> {
    fn max_pool2d(
        x: FloatTensor<Self>,
        kernel_size: [usize;
    fn max_pool1d_with_indices(
        x: FloatTensor<Self>,
        kernel_size: usize,
        stride: usize,
        padding: usize,
        dilation: usize,
    ) -> MaxPool1dWithIndices<Self> {
    fn max_pool2d_with_indices(
        x: FloatTensor<Self>,
        kernel_size: [usize;
    fn max_pool1d_with_indices_backward(
        x: FloatTensor<Self>,
        kernel_size: usize,
        stride: usize,
        padding: usize,
        dilation: usize,
        output_grad: FloatTensor<Self>,
        indices: IntTensor<Self>,
    ) -> MaxPool1dBackward<Self> {
    fn max_pool2d_with_indices_backward(
        x: FloatTensor<Self>,
        kernel_size: [usize;
    fn adaptive_avg_pool1d(x: FloatTensor<Self>, output_size: usize) -> FloatTensor<Self> {
    fn adaptive_avg_pool2d(x: FloatTensor<Self>, output_size: [usize;
    fn adaptive_avg_pool1d_backward(
        x: FloatTensor<Self>,
        grad: FloatTensor<Self>,
    ) -> FloatTensor<Self> {
    fn adaptive_avg_pool2d_backward(
        x: FloatTensor<Self>,
        grad: FloatTensor<Self>,
    ) -> FloatTensor<Self> {
    fn interpolate(
        x: FloatTensor<Self>,
        output_size: [usize;
    fn interpolate_backward(
        x: FloatTensor<Self>,
        grad: FloatTensor<Self>,
        output_size: [usize;
    fn deform_conv2d(
        x: FloatTensor<Self>,
        offset: FloatTensor<Self>,
        weight: FloatTensor<Self>,
        mask: Option<FloatTensor<Self>>,
        bias: Option<FloatTensor<Self>>,
        options: DeformConvOptions<2>,
    ) -> FloatTensor<Self> {
    fn deform_conv2d_backward(
        x: FloatTensor<Self>,
        offset: FloatTensor<Self>,
        weight: FloatTensor<Self>,
        mask: Option<FloatTensor<Self>>,
        bias: Option<FloatTensor<Self>>,
        output_grad: FloatTensor<Self>,
        options: DeformConvOptions<2>,
    ) -> DeformConv2dBackward<Self> {
impl<R: RunnerChannel> QTensorOps<Self> for BackendRouter<R> {
    fn q_from_data(_data: TensorData, _device: &Device<Self>) -> QuantizedTensor<Self> {
    fn quantize(
        _tensor: FloatTensor<Self>,
        _scheme: &QuantScheme,
        _qparams: QuantizationParametersPrimitive<Self>,
    ) -> QuantizedTensor<Self> {
    fn quantize_dynamic(
        _tensor: FloatTensor<Self>,
        _scheme: &QuantScheme,
    ) -> QuantizedTensor<Self> {
    fn dequantize(_tensor: QuantizedTensor<Self>) -> FloatTensor<Self> {
    fn q_device(_tensor: &QuantizedTensor<Self>) -> Device<Self> {
    fn q_to_device(
        _tensor: QuantizedTensor<Self>,
        _device: &Device<Self>,
    ) -> QuantizedTensor<Self> {
    fn q_reshape(_tensor: QuantizedTensor<Self>, _shape: Shape) -> QuantizedTensor<Self> {
    async fn q_into_data(_tensor: QuantizedTensor<Self>) -> TensorData {
    fn q_swap_dims(
        _tensor: QuantizedTensor<Self>,
        _dim1: usize,
        _dim2: usize,
    ) -> QuantizedTensor<Self> {
    fn q_permute(_tensor: QuantizedTensor<Self>, _axes: &[usize]) -> QuantizedTensor<Self> {
    fn q_flip(_tensor: QuantizedTensor<Self>, _axes: &[usize]) -> QuantizedTensor<Self> {
    fn q_gather(
        _dim: usize,
        _tensor: QuantizedTensor<Self>,
        _indices: IntTensor<Self>,
    ) -> QuantizedTensor<Self> {
    fn q_select(
        _tensor: QuantizedTensor<Self>,
        _dim: usize,
        _indices: IntTensor<Self>,
    ) -> QuantizedTensor<Self> {
    fn q_slice(_tensor: QuantizedTensor<Self>, _ranges: &[Range<usize>]) -> QuantizedTensor<Self> {
    fn q_expand(_tensor: QuantizedTensor<Self>, _shape: Shape) -> QuantizedTensor<Self> {
impl<R: RunnerChannel> TransactionOps<Self> for BackendRouter<R> {
/// A runner's context contains a [handle container](HandleContainer) to manage
/// (i.e., fetch and update) existing tensors.
pub struct RunnerContext<B: BackendIr> {
/// Handle container to retrieve tensors based on their intermediate representation.
impl<B: BackendIr> RunnerContext<B> {
/// Create a new (uninitialized) empty tensor and returns its corresponding [tensor id](TensorId).
    fn create_empty_handle(&mut self) -> TensorId {
/// A runner is responsible for executing tensor operations for a given [intermediate backend](BackendIr).
pub struct Runner<B: BackendIr> {
// Mutex for the mutable handles
impl<B: BackendIr> Runner<B> {
/// Create a new runner.
    pub fn new(device: B::Device) -> Self {
/// Get the tensor handle for the given [tensor representation](TensorIr).
    pub(crate) fn get_tensor_handle(&self, tensor: &TensorIr) -> B::Handle {
/// Create a tensor with the given handle and shape.
    pub(crate) fn register_tensor<C: RunnerClient>(
        &self,
        handle: B::Handle,
        shape: Vec<usize>,
        dtype: DType,
        client: C,
    ) -> RouterTensor<C> {
/// Register a tensor from its data and id.
    pub fn register_tensor_data_id(&self, id: TensorId, data: TensorData) {
/// Register a tensor and returns its intermediate representation.
    pub fn register_tensor_data_desc(&self, data: TensorData) -> TensorIr {
/// Register an empty tensor and returns its intermediate representation.
    pub fn register_empty_tensor_desc(&self, shape: Vec<usize>, dtype: DType) -> TensorIr {
    pub(crate) fn register_float_tensor_desc(
        &self,
        shape: Vec<usize>,
        dtype: FloatDType,
    ) -> TensorIr {
impl<B: BackendIr> RunnerClient for Runner<B> {
    type Device = B::Device;
/// Execute a tensor operation.
    fn register(&self, op: OperationIr) {
// Remove unused tensor handles
// For every op: get the input(s), execute the operation and register the output(s)
// Nothing to do.
    fn read_tensor(&self, tensor: TensorIr) -> DynFut<TensorData> {
        enum Output<B: Backend> {
    fn register_tensor_data(&self, data: TensorData) -> RouterTensor<Self> {
    fn register_empty_tensor(&self, shape: Vec<usize>, dtype: DType) -> RouterTensor<Self> {
    fn register_float_tensor(&self, shape: Vec<usize>, dtype: FloatDType) -> RouterTensor<Self> {
    fn device(&self) -> Self::Device {
    fn sync(&self) {
    fn seed(&self, seed: u64) {
/// Tensor primitive for the [router backend](crate::BackendRouter).
pub struct RouterTensor<C: RunnerClient> {
/// The client that has this tensor
impl<C: RunnerClient> TensorMetadata for RouterTensor<C> {
    fn dtype(&self) -> DType {
    fn shape(&self) -> Shape {
impl<C: RunnerClient> RouterTensor<C> {
/// Create a new router tensor.
    pub fn new(id: TensorId, shape: Vec<usize>, dtype: DType, client: C) -> Self {
    pub(crate) async fn into_data(self) -> TensorData {
/// Get the ir for this tensor
    pub fn into_ir(mut self) -> TensorIr {
// Avoids an unwanted drop on the same thread.
//
// Since `drop` is called after `into_ir`, we must not register a drop if the tensor
// was consumed with a `ReadWrite` status.
    pub(crate) fn to_ir_out(&self) -> TensorIr {
    pub(crate) fn status(&self, count: u32) -> TensorStatus {
impl<C: RunnerClient> core::fmt::Debug for RouterTensor<C> {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
impl<C: RunnerClient> Clone for RouterTensor<C> {
    fn clone(&self) -> Self {
impl<C: RunnerClient> Drop for RouterTensor<C> {
    fn drop(&mut self) {
/// Implement multi backend types, with enums having one variant per backend.
// Match the default backend and at least one other backend, with rest being optional
/// Module containing the essential types for multi-backend operations.
///
/// - `Handle`: the type used to point to a tensor (defined for all backends).
/// - `MultiRunnerClient`: a client for multiple runners (each responsible to execute tensor operations on a given backend).
/// - `DirectChannel`: a local channel with direct connection to the backend runner clients.
/// - `ByteBridge`: a simple multi-backend bridge that transfers tensors via the underlying [tensor data](burn_tensor::TensorData).
///
/// Each enum type is defined with backend identifiers as variant names (e.g., `B1` and `B2` for dual backends).
/// The type that can be used to point to a tensor of any kind.
/// Each backend has its own variant.
            pub enum Handle<$DefaultBackend: BackendIr, $($OtherBackend: BackendIr),+> {
/// The device type used by a backend.
/// Each backend has its own variant.
            pub enum MultiDevice<$DefaultBackend: Backend, $($OtherBackend: Backend),+> {
            impl<$DefaultBackend: Backend, $($OtherBackend: Backend),+> PartialEq for MultiDevice<$DefaultBackend, $($OtherBackend),+> {
                fn eq(&self, other: &Self) -> bool {
// Default implementation always returns the first backend's device
            impl<$DefaultBackend: Backend, $($OtherBackend: Backend),+> Default for MultiDevice<$DefaultBackend, $($OtherBackend),+> {
                fn default() -> Self {
            impl<$DefaultBackend: Backend, $($OtherBackend: Backend),+> DeviceOps for MultiDevice<$DefaultBackend, $($OtherBackend),+> {
                fn id(&self) -> DeviceId {
/// A local client with multiple runners (each responsible to execute tensor operations on a given backend).
            pub enum MultiRunnerClient<$DefaultBackend: BackendIr, $($OtherBackend: BackendIr),+> {
            impl<$DefaultBackend: BackendIr, $($OtherBackend: BackendIr),+> RunnerClient for MultiRunnerClient<$DefaultBackend, $($OtherBackend),+>
            {
               type Device = MultiDevice<$DefaultBackend, $($OtherBackend),+>;
                fn register(&self, op: OperationIr) {
                fn read_tensor(&self, tensor: TensorIr) -> DynFut<TensorData> {
                fn register_tensor_data(&self, data: TensorData) -> RouterTensor<Self> {
                fn register_empty_tensor(&self, shape: Vec<usize>, dtype: DType) -> RouterTensor<Self> {
                fn register_float_tensor(&self, shape: Vec<usize>, dtype: burn_tensor::FloatDType) -> RouterTensor<Self> {
                fn device(&self) -> Self::Device {
                fn sync(&self) {
                fn seed(&self, seed: u64) {
            impl<$DefaultBackend: BackendIr, $($OtherBackend: BackendIr),+, Br> RunnerChannel for DirectChannel<($DefaultBackend, $($OtherBackend),+), Br>
            where
                Br: MultiBackendBridge<TensorHandle = Handle<$DefaultBackend, $($OtherBackend),+>, Device = MultiDevice<$DefaultBackend, $($OtherBackend),+>>,
            {
                type Device = Br::Device;
                type Bridge = Br;
                type FloatElem = $DefaultBackend::FloatElem;
                type IntElem = $DefaultBackend::IntElem;
                type BoolElem = $DefaultBackend::BoolElem;
                type Client = MultiRunnerClient<$DefaultBackend, $($OtherBackend),+>;
                fn init_client(device: &Self::Device) -> Self::Client {
                fn get_tensor_handle(
                    tensor: &TensorIr,
                    client: &Self::Client,
                ) -> <Self::Bridge as MultiBackendBridge>::TensorHandle {
                fn register_tensor(
                    client: &Self::Client,
                    handle: <Self::Bridge as MultiBackendBridge>::TensorHandle,
                    shape: Vec<usize>,
                    dtype: DType,
                ) -> RouterTensor<Self::Client> {
                fn name(_device: &Self::Device) -> String {
            impl<$DefaultBackend: BackendIr, $($OtherBackend: BackendIr),+> MultiBackendBridge for ByteBridge<($DefaultBackend, $($OtherBackend),+)> {
                type TensorHandle = Handle<$DefaultBackend, $($OtherBackend),+>;
                type Device = MultiDevice<$DefaultBackend, $($OtherBackend),+>;
                fn change_backend_float(
                    tensor: Self::TensorHandle,
                    shape: Shape,
                    target_device: &Self::Device,
                ) -> Self::TensorHandle {
                fn change_backend_int(
                    tensor: Self::TensorHandle,
                    shape: Shape,
                    target_device: &Self::Device,
                ) -> Self::TensorHandle {
                fn change_backend_bool(
                    tensor: Self::TensorHandle,
                    shape: Shape,
                    target_device: &Self::Device,
                ) -> Self::TensorHandle {
// Bridge for the same backend
// Byte bridge between two backends
// Implement multi-backend types and byte bridge for up to 4 backends
// cannot find a wgpu adapter on windows CI
    fn should_support_dual_byte_bridge() {
// The LIBTORCH environment variable can be used to specify the directory
// where libtorch has been installed.
// When not specified this script downloads the cpu version for libtorch
// and extracts it in OUT_DIR.
//
// On Linux, the TORCH_CUDA_VERSION environment variable can be used,
// like 9.0, 90, or cu90 to specify the version of CUDA to use for libtorch.
const PYTHON_PRINT_PYTORCH_DETAILS: &str = r"
import torch
from torch.utils import cpp_extension
print('LIBTORCH_VERSION:', torch.__version__.split('+')[0])
print('LIBTORCH_CXX11:', torch._C._GLIBCXX_USE_CXX11_ABI)
for include_path in cpp_extension.include_paths():
  print('LIBTORCH_INCLUDE:', include_path)
for library_path in cpp_extension.library_paths():
  print('LIBTORCH_LIB:', library_path)
";
/// The device struct when using the `tch` backend.
///
/// Note that you need to provide the device index when using Cuda.
///
/// # Example
///
/// ```no_run
/// use burn_tch::LibTorchDevice;
///
/// let device_gpu_1 = LibTorchDevice::Cuda(0); // First GPU
/// let device_gpu_2 = LibTorchDevice::Cuda(1); // Second GPU
/// let device_cpu = LibTorchDevice::Cpu; // CPU
/// let device_mps = LibTorchDevice::Mps; // Metal Performance Shaders
/// let device_vulkan = LibTorchDevice::Vulkan; // Vulkan
/// ```
pub enum LibTorchDevice {
/// CPU device.
/// Cuda device with the given index. The index is the index of the Cuda device in the list of
/// all Cuda devices found on the system.
/// Metal Performance Shaders device.
/// Vulkan device.
impl From<LibTorchDevice> for tch::Device {
    fn from(device: LibTorchDevice) -> Self {
impl From<tch::Device> for LibTorchDevice {
    fn from(device: tch::Device) -> Self {
impl DeviceOps for LibTorchDevice {
    fn id(&self) -> burn_tensor::backend::DeviceId {
impl Default for LibTorchDevice {
    fn default() -> Self {
/// Tensor backend that uses `LibTorch` with the [tch] crate for executing tensor operations.
///
/// This backend is compatible with a wide range of hardwares ranging from CPUs to GPUs, but
/// requires `LibTorch` to be installed correctly. The CPU version can be downloaded
/// automatically and the CUDA version as well by setting the `TORCH_CUDA_VERSION` environment
/// variable. For more complex configurations, check out the manual installation for
/// [burn-tch](https://github.com/tracel-ai/burn/tree/main/crates/burn-tch).
///
/// Refer to the [tch] crate for more information.
pub struct LibTorch<E = f32, Q = i8> {
impl<E: TchElement, Q: QuantElement> Backend for LibTorch<E, Q> {
    type Device = LibTorchDevice;
    type FloatTensorPrimitive = TchTensor;
    type FloatElem = E;
    type IntTensorPrimitive = TchTensor;
    type IntElem = i64;
    type BoolTensorPrimitive = TchTensor;
    type BoolElem = bool;
    type QuantizedTensorPrimitive = TchQTensor;
    type QuantizedEncoding = Q;
    fn seed(seed: u64) {
    fn ad_enabled() -> bool {
    fn name(device: &Self::Device) -> String {
    fn sync(device: &Self::Device) {
// When there is no explicit way to synchronize, we write and read one value to sync
fn main() {
// Creation of two tensors, the first with explicit values and the second one with ones, with the same shape as the first
// Print the element-wise addition of the two tensors.
fn main() {
// Creation of two tensors, the first with explicit values and the second one with ones, with the same shape as the first
// Print the element-wise addition of the two tensors.
fn main() {
// Creation of two tensors, the first with explicit values and the second one with ones, with the same shape as the first
// Print the element-wise addition of the two tensors.
/// The element type for the tch backend.
pub trait TchElement: Element + tch::kind::Element {
impl TchElement for f64 {
impl TchElement for f32 {
impl TchElement for f16 {
impl TchElement for bf16 {
impl TchElement for i64 {
impl TchElement for i32 {
impl TchElement for i16 {
impl TchElement for i8 {
impl TchElement for u8 {
impl TchElement for bool {
/// A quantized element for the tch backend.
pub trait QuantElement: TchElement {
impl QuantElement for i8 {
//! Burn Tch Backend
    type TestBackend = crate::LibTorch<f32>;
    type TestTensor<const D: usize> = burn_tensor::Tensor<TestBackend, D>;
    type TestTensorInt<const D: usize> = burn_tensor::Tensor<TestBackend, D, burn_tensor::Int>;
    type TestTensorBool<const D: usize> = burn_tensor::Tensor<TestBackend, D, burn_tensor::Bool>;
impl<E: TchElement, Q: QuantElement> ActivationOps<Self> for LibTorch<E, Q> {
    fn relu(tensor: TchTensor) -> TchTensor {
    fn gelu(tensor: TchTensor) -> TchTensor {
    fn gelu_backward(tensor: TchTensor, grad: TchTensor) -> TchTensor {
    fn sigmoid(tensor: TchTensor) -> TchTensor {
    fn log_sigmoid(tensor: TchTensor) -> TchTensor {
// NOTE: we don't override log_sigmoid_backward because Torch has a special backward
// formula that uses a buffer with computed values from the forward pass
// no in-place log_sigmoid_
pub struct TchOps {
// e: PhantomData<E>,
impl TchOps {
    pub fn to_device(tensor: TchTensor, device: &LibTorchDevice) -> TchTensor {
// We have to manually check if the device is the same, since when it's the case, we need to keep
// the same storage reference and not create a new one.
    pub fn reshape(tensor: TchTensor, shape: Shape) -> TchTensor {
    pub fn repeat_dim(tensor: TchTensor, dim: usize, times: usize) -> TchTensor {
    pub fn slice(tensor: TchTensor, ranges: &[Range<usize>]) -> TchTensor {
    pub fn slice_assign(tensor: TchTensor, ranges: &[Range<usize>], value: TchTensor) -> TchTensor {
// Copy the input tensor if we can't mutate it.
    pub fn gather(dim: usize, tensor: TchTensor, indices: TchTensor) -> TchTensor {
    pub fn scatter(
        dim: usize,
        tensor: TchTensor,
        indices: TchTensor,
        value: TchTensor,
    ) -> TchTensor {
    pub fn index_select_dim(tensor: TchTensor, dim: usize, indices: TchTensor) -> TchTensor {
    pub fn select_assign(
        tensor: TchTensor,
        dim: usize,
        indices: TchTensor,
        value: TchTensor,
    ) -> TchTensor {
    pub fn cat(tensors: Vec<TchTensor>, dim: usize) -> TchTensor {
    pub fn equal(lhs: TchTensor, rhs: TchTensor) -> TchTensor {
    pub fn equal_elem<S: Into<tch::Scalar> + Clone>(lhs: TchTensor, rhs: S) -> TchTensor {
    pub fn greater(lhs: TchTensor, rhs: TchTensor) -> TchTensor {
    pub fn greater_elem<S: Into<tch::Scalar> + Clone>(lhs: TchTensor, rhs: S) -> TchTensor {
    pub fn greater_equal(lhs: TchTensor, rhs: TchTensor) -> TchTensor {
    pub fn greater_equal_elem<S: Into<Scalar> + Clone>(lhs: TchTensor, rhs: S) -> TchTensor {
    pub fn lower(lhs: TchTensor, rhs: TchTensor) -> TchTensor {
    pub fn lower_elem<S: Into<Scalar> + Clone>(lhs: TchTensor, rhs: S) -> TchTensor {
    pub fn lower_equal(lhs: TchTensor, rhs: TchTensor) -> TchTensor {
    pub fn lower_equal_elem<S: Into<Scalar> + Clone>(lhs: TchTensor, rhs: S) -> TchTensor {
    pub fn add(lhs: TchTensor, rhs: TchTensor) -> TchTensor {
    pub fn sub(lhs: TchTensor, rhs: TchTensor) -> TchTensor {
    pub fn mul(lhs: TchTensor, rhs: TchTensor) -> TchTensor {
    pub fn div(lhs: TchTensor, rhs: TchTensor) -> TchTensor {
    pub fn remainder(lhs: TchTensor, rhs: TchTensor) -> TchTensor {
    pub fn mean(tensor: TchTensor) -> TchTensor {
// view as 1d tensor
    pub fn mean_dim(tensor: TchTensor, dim: usize) -> TchTensor {
    pub fn sum(tensor: TchTensor) -> TchTensor {
// view as 1d tensor
    pub fn sum_dim(tensor: TchTensor, dim: usize) -> TchTensor {
    pub fn prod(tensor: TchTensor) -> TchTensor {
// view as 1d tensor
    pub fn prod_dim(tensor: TchTensor, dim: usize) -> TchTensor {
    pub fn argmax(tensor: TchTensor, dim: usize) -> TchTensor {
    pub fn argmin(tensor: TchTensor, dim: usize) -> TchTensor {
    pub fn max_dim(tensor: TchTensor, dim: usize) -> TchTensor {
    pub fn max_dim_with_indices(tensor: TchTensor, dim: usize) -> (TchTensor, TchTensor) {
    pub fn min_dim(tensor: TchTensor, dim: usize) -> TchTensor {
    pub fn min_dim_with_indices(tensor: TchTensor, dim: usize) -> (TchTensor, TchTensor) {
    pub fn clamp_min<S: Into<tch::Scalar> + Clone + Copy>(tensor: TchTensor, min: S) -> TchTensor {
    pub fn clamp_max<S: Into<tch::Scalar> + Clone + Copy>(tensor: TchTensor, max: S) -> TchTensor {
    pub fn clamp<S: Into<tch::Scalar> + Clone + Copy>(
        tensor: TchTensor,
        min: S,
        max: S,
    ) -> TchTensor {
    pub fn swap_dims(tensor: TchTensor, dim1: usize, dim2: usize) -> TchTensor {
    pub fn permute(tensor: TchTensor, axes: &[usize]) -> TchTensor {
    pub fn flip(tensor: TchTensor, axes: &[usize]) -> TchTensor {
    pub fn powf(tensor: TchTensor, exponent: TchTensor) -> TchTensor {
    pub fn sign(tensor: TchTensor) -> TchTensor {
    pub fn expand(tensor: TchTensor, shape: Shape) -> TchTensor {
    pub fn sort(tensor: TchTensor, dim: usize, descending: bool) -> TchTensor {
    pub fn sort_with_indices(
        tensor: TchTensor,
        dim: usize,
        descending: bool,
    ) -> (TchTensor, TchTensor) {
    pub fn argsort(tensor: TchTensor, dim: usize, descending: bool) -> TchTensor {
    pub fn bitwise_and(lhs: TchTensor, rhs: TchTensor) -> TchTensor {
    pub fn bitwise_and_scalar<S: Into<Scalar> + Clone>(tensor: TchTensor, scalar: S) -> TchTensor {
    pub fn bitwise_or(lhs: TchTensor, rhs: TchTensor) -> TchTensor {
    pub fn bitwise_or_scalar<S: Into<Scalar> + Clone>(tensor: TchTensor, scalar: S) -> TchTensor {
    pub fn bitwise_xor(lhs: TchTensor, rhs: TchTensor) -> TchTensor {
    pub fn bitwise_xor_scalar<S: Into<Scalar> + Clone>(tensor: TchTensor, scalar: S) -> TchTensor {
    pub fn bitwise_not(tensor: TchTensor) -> TchTensor {
    pub fn bitwise_left_shift(lhs: TchTensor, rhs: TchTensor) -> TchTensor {
    pub fn bitwise_left_shift_scalar<S: Into<Scalar> + Clone>(
        tensor: TchTensor,
        scalar: S,
    ) -> TchTensor {
    pub fn bitwise_right_shift(lhs: TchTensor, rhs: TchTensor) -> TchTensor {
    pub fn bitwise_right_shift_scalar<S: Into<Scalar> + Clone>(
        tensor: TchTensor,
        scalar: S,
    ) -> TchTensor {
impl<E: TchElement, Q: QuantElement> BoolTensorOps<Self> for LibTorch<E, Q> {
    fn bool_from_data(data: TensorData, device: &LibTorchDevice) -> TchTensor {
    fn bool_repeat_dim(tensor: TchTensor, dim: usize, times: usize) -> TchTensor {
    async fn bool_into_data(tensor: TchTensor) -> TensorData {
    fn bool_to_device(tensor: TchTensor, device: &LibTorchDevice) -> TchTensor {
    fn bool_reshape(tensor: TchTensor, shape: Shape) -> TchTensor {
    fn bool_device(tensor: &TchTensor) -> LibTorchDevice {
    fn bool_empty(shape: Shape, device: &<LibTorch<E> as Backend>::Device) -> TchTensor {
    fn bool_slice(tensor: TchTensor, ranges: &[Range<usize>]) -> TchTensor {
    fn bool_slice_assign(
        tensor: TchTensor,
        ranges: &[Range<usize>],
        value: TchTensor,
    ) -> TchTensor {
    fn bool_cat(tensors: Vec<TchTensor>, dim: usize) -> TchTensor {
    fn bool_equal(lhs: TchTensor, rhs: TchTensor) -> TchTensor {
    fn bool_not(tensor: TchTensor) -> TchTensor {
    fn bool_and(lhs: TchTensor, rhs: TchTensor) -> TchTensor {
    fn bool_or(lhs: TchTensor, rhs: TchTensor) -> TchTensor {
    fn bool_into_int(tensor: TchTensor) -> TchTensor {
    fn bool_into_float(tensor: TchTensor) -> TchTensor {
    fn bool_swap_dims(tensor: TchTensor, dim1: usize, dim2: usize) -> TchTensor {
    fn bool_permute(tensor: TchTensor, axes: &[usize]) -> TchTensor {
    fn bool_flip(tensor: TchTensor, axes: &[usize]) -> TchTensor {
    async fn bool_argwhere(tensor: TchTensor) -> TchTensor {
    fn bool_expand(tensor: TchTensor, shape: Shape) -> TchTensor {
impl<E: TchElement, Q: QuantElement> IntTensorOps<Self> for LibTorch<E, Q> {
    fn int_from_data(data: TensorData, device: &LibTorchDevice) -> TchTensor {
    fn int_repeat_dim(tensor: TchTensor, dim: usize, times: usize) -> TchTensor {
    async fn int_into_data(tensor: TchTensor) -> TensorData {
    fn int_to_device(tensor: TchTensor, device: &LibTorchDevice) -> TchTensor {
    fn int_reshape(tensor: TchTensor, shape: Shape) -> TchTensor {
    fn int_device(tensor: &TchTensor) -> LibTorchDevice {
    fn int_empty(shape: Shape, device: &<LibTorch<E> as Backend>::Device) -> TchTensor {
    fn int_slice(tensor: TchTensor, ranges: &[Range<usize>]) -> TchTensor {
    fn int_slice_assign(tensor: TchTensor, ranges: &[Range<usize>], value: TchTensor) -> TchTensor {
    fn int_cat(tensors: Vec<TchTensor>, dim: usize) -> TchTensor {
    fn int_equal(lhs: TchTensor, rhs: TchTensor) -> TchTensor {
    fn int_equal_elem(lhs: TchTensor, rhs: i64) -> TchTensor {
    fn int_greater(lhs: TchTensor, rhs: TchTensor) -> TchTensor {
    fn int_greater_elem(lhs: TchTensor, rhs: i64) -> TchTensor {
    fn int_greater_equal(lhs: TchTensor, rhs: TchTensor) -> TchTensor {
    fn int_greater_equal_elem(lhs: TchTensor, rhs: i64) -> TchTensor {
    fn int_lower(lhs: TchTensor, rhs: TchTensor) -> TchTensor {
    fn int_lower_elem(lhs: TchTensor, rhs: i64) -> TchTensor {
    fn int_lower_equal(lhs: TchTensor, rhs: TchTensor) -> TchTensor {
    fn int_lower_equal_elem(lhs: TchTensor, rhs: i64) -> TchTensor {
    fn int_add(lhs: TchTensor, rhs: TchTensor) -> TchTensor {
    fn int_add_scalar(lhs: TchTensor, rhs: i64) -> TchTensor {
    fn int_sub(lhs: TchTensor, rhs: TchTensor) -> TchTensor {
    fn int_sub_scalar(lhs: TchTensor, rhs: i64) -> TchTensor {
    fn int_mul(lhs: TchTensor, rhs: TchTensor) -> TchTensor {
    fn int_mul_scalar(lhs: TchTensor, rhs: i64) -> TchTensor {
    fn int_div(lhs: TchTensor, rhs: TchTensor) -> TchTensor {
    fn int_div_scalar(lhs: TchTensor, rhs: i64) -> TchTensor {
    fn int_remainder(lhs: TchTensor, rhs: TchTensor) -> TchTensor {
    fn int_remainder_scalar(lhs: TchTensor, rhs: i64) -> TchTensor {
    fn int_neg(tensor: TchTensor) -> TchTensor {
    fn int_zeros(shape: Shape, device: &<LibTorch<E> as Backend>::Device) -> TchTensor {
    fn int_ones(shape: Shape, device: &<LibTorch<E> as Backend>::Device) -> TchTensor {
    fn int_full(
        shape: Shape,
        fill_value: i64,
        device: &<LibTorch<E> as Backend>::Device,
    ) -> TchTensor {
    fn int_sum(tensor: TchTensor) -> TchTensor {
    fn int_sum_dim(tensor: TchTensor, dim: usize) -> TchTensor {
    fn int_prod(tensor: TchTensor) -> TchTensor {
    fn int_prod_dim(tensor: TchTensor, dim: usize) -> TchTensor {
    fn int_mean(tensor: TchTensor) -> TchTensor {
    fn int_mean_dim(tensor: TchTensor, dim: usize) -> TchTensor {
    fn int_gather(dim: usize, tensor: TchTensor, indices: TchTensor) -> TchTensor {
    fn int_scatter(
        dim: usize,
        tensor: TchTensor,
        indices: TchTensor,
        value: TchTensor,
    ) -> TchTensor {
    fn int_select(tensor: TchTensor, dim: usize, indices: TchTensor) -> TchTensor {
    fn int_select_assign(
        tensor: TchTensor,
        dim: usize,
        indices: TchTensor,
        value: TchTensor,
    ) -> TchTensor {
    fn int_mask_where(tensor: TchTensor, mask: TchTensor, source: TchTensor) -> TchTensor {
    fn int_mask_fill(tensor: TchTensor, mask: TchTensor, value: i64) -> TchTensor {
    fn int_argmax(tensor: TchTensor, dim: usize) -> TchTensor {
    fn int_argmin(tensor: TchTensor, dim: usize) -> TchTensor {
    fn int_max_dim(tensor: TchTensor, dim: usize) -> TchTensor {
    fn int_max_dim_with_indices(tensor: TchTensor, dim: usize) -> (TchTensor, TchTensor) {
    fn int_min_dim(tensor: TchTensor, dim: usize) -> TchTensor {
    fn int_min_dim_with_indices(tensor: TchTensor, dim: usize) -> (TchTensor, TchTensor) {
    fn int_clamp_min(tensor: TchTensor, min: i64) -> TchTensor {
    fn int_clamp_max(tensor: TchTensor, max: i64) -> TchTensor {
    fn int_clamp(tensor: TchTensor, min: i64, max: i64) -> TchTensor {
    fn int_abs(tensor: TchTensor) -> TchTensor {
    fn int_into_float(tensor: TchTensor) -> TchTensor {
    fn int_swap_dims(tensor: IntTensor<Self>, dim1: usize, dim2: usize) -> IntTensor<Self> {
    fn int_random(shape: Shape, distribution: Distribution, device: &LibTorchDevice) -> TchTensor {
    fn int_arange(range: Range<i64>, device: &LibTorchDevice) -> TchTensor {
    fn int_permute(tensor: IntTensor<Self>, axes: &[usize]) -> IntTensor<Self> {
    fn int_flip(tensor: IntTensor<Self>, axes: &[usize]) -> IntTensor<Self> {
    fn int_sign(tensor: IntTensor<Self>) -> IntTensor<Self> {
    fn int_expand(tensor: IntTensor<Self>, shape: Shape) -> IntTensor<Self> {
    fn int_sort(tensor: IntTensor<Self>, dim: usize, descending: bool) -> IntTensor<Self> {
    fn int_argsort(tensor: IntTensor<Self>, dim: usize, descending: bool) -> IntTensor<Self> {
    fn bitwise_and(lhs: IntTensor<Self>, rhs: IntTensor<Self>) -> IntTensor<Self> {
    fn bitwise_or(lhs: IntTensor<Self>, rhs: IntTensor<Self>) -> IntTensor<Self> {
    fn bitwise_xor(lhs: IntTensor<Self>, rhs: IntTensor<Self>) -> IntTensor<Self> {
    fn bitwise_not(tensor: IntTensor<Self>) -> IntTensor<Self> {
    fn bitwise_and_scalar(
        lhs: IntTensor<Self>,
        rhs: burn_tensor::ops::IntElem<Self>,
    ) -> IntTensor<Self> {
    fn bitwise_or_scalar(
        lhs: IntTensor<Self>,
        rhs: burn_tensor::ops::IntElem<Self>,
    ) -> IntTensor<Self> {
    fn bitwise_xor_scalar(
        lhs: IntTensor<Self>,
        rhs: burn_tensor::ops::IntElem<Self>,
    ) -> IntTensor<Self> {
    fn bitwise_left_shift(lhs: IntTensor<Self>, rhs: IntTensor<Self>) -> IntTensor<Self> {
    fn bitwise_right_shift(lhs: IntTensor<Self>, rhs: IntTensor<Self>) -> IntTensor<Self> {
    fn bitwise_left_shift_scalar(
        lhs: IntTensor<Self>,
        rhs: burn_tensor::ops::IntElem<Self>,
    ) -> IntTensor<Self> {
    fn bitwise_right_shift_scalar(
        lhs: IntTensor<Self>,
        rhs: burn_tensor::ops::IntElem<Self>,
    ) -> IntTensor<Self> {
impl<E: TchElement, Q: QuantElement> ModuleOps<Self> for LibTorch<E, Q> {
    fn embedding(weights: TchTensor, indices: TchTensor) -> TchTensor {
    fn embedding_backward(weights: TchTensor, output: TchTensor, indices: TchTensor) -> TchTensor {
    fn conv1d(
        x: TchTensor,
        weight: TchTensor,
        bias: Option<TchTensor>,
        options: ConvOptions<1>,
    ) -> TchTensor {
    fn conv2d(
        x: TchTensor,
        weight: TchTensor,
        bias: Option<TchTensor>,
        options: ConvOptions<2>,
    ) -> TchTensor {
    fn conv3d(
        x: TchTensor,
        weight: TchTensor,
        bias: Option<TchTensor>,
        options: ConvOptions<3>,
    ) -> TchTensor {
    fn deform_conv2d(
        _x: TchTensor,
        _offset: TchTensor,
        _weight: TchTensor,
        _mask: Option<TchTensor>,
        _bias: Option<TchTensor>,
        _options: DeformConvOptions<2>,
    ) -> TchTensor {
    fn deform_conv2d_backward(
        _x: TchTensor,
        _offset: TchTensor,
        _weight: TchTensor,
        _mask: Option<TchTensor>,
        _bias: Option<TchTensor>,
        _out_grad: TchTensor,
        _options: DeformConvOptions<2>,
    ) -> DeformConv2dBackward<Self> {
    fn conv_transpose1d(
        x: TchTensor,
        weight: TchTensor,
        bias: Option<TchTensor>,
        options: ConvTransposeOptions<1>,
    ) -> TchTensor {
    fn conv_transpose2d(
        x: TchTensor,
        weight: TchTensor,
        bias: Option<TchTensor>,
        options: ConvTransposeOptions<2>,
    ) -> TchTensor {
    fn conv_transpose3d(
        x: TchTensor,
        weight: TchTensor,
        bias: Option<TchTensor>,
        options: ConvTransposeOptions<3>,
    ) -> TchTensor {
    fn avg_pool1d(
        x: TchTensor,
        kernel_size: usize,
        stride: usize,
        padding: usize,
        count_include_pad: bool,
    ) -> TchTensor {
    fn avg_pool2d(
        x: TchTensor,
        kernel_size: [usize;
    fn avg_pool2d_backward(
        x: TchTensor,
        grad: TchTensor,
        kernel_size: [usize;
    fn max_pool1d(
        x: TchTensor,
        kernel_size: usize,
        stride: usize,
        padding: usize,
        dilation: usize,
    ) -> TchTensor {
    fn max_pool1d_with_indices(
        x: TchTensor,
        kernel_size: usize,
        stride: usize,
        padding: usize,
        dilation: usize,
    ) -> MaxPool1dWithIndices<LibTorch<E, Q>> {
    fn max_pool2d(
        x: TchTensor,
        kernel_size: [usize;
    fn max_pool2d_with_indices(
        x: TchTensor,
        kernel_size: [usize;
    fn max_pool2d_with_indices_backward(
        x: TchTensor,
        kernel_size: [usize;
    fn adaptive_avg_pool2d(x: TchTensor, output_size: [usize;
    fn adaptive_avg_pool2d_backward(x: TchTensor, grad: TchTensor) -> TchTensor {
    fn adaptive_avg_pool1d(x: TchTensor, output_size: usize) -> TchTensor {
    fn interpolate(
        x: TchTensor,
        output_size: [usize;
    fn interpolate_backward(
        x: TchTensor,
        grad: TchTensor,
        output_size: [usize;
fn quantize<E: TchElement, Q: QuantElement>(
    tensor: tch::Tensor,
    scheme: &QuantScheme,
    qparams: &QParams<E, Q>,
) -> tch::Tensor {
// Quantize only works on Float Tensor
impl<E: TchElement, Q: QuantElement> QTensorOps<Self> for LibTorch<E, Q> {
    fn q_from_data(data: TensorData, device: &LibTorchDevice) -> QuantizedTensor<Self> {
// NOTE: tch-rs doesn't have `from_blob_quantized_*` APIs
// https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/quantized/Quantizer.cpp#L322
// So for now we have to load the dequantized values to quantize them back since the dequantization
// methods take the values provided when quantizing.
    fn quantize(
        tensor: FloatTensor<Self>,
        scheme: &QuantScheme,
        qparams: QuantizationParametersPrimitive<Self>,
    ) -> QuantizedTensor<Self> {
// Quantize only works on Float Tensor
    fn quantize_dynamic(tensor: FloatTensor<Self>, scheme: &QuantScheme) -> QuantizedTensor<Self> {
/*reduce_range*/ false)
            }
        };
        TchQTensor {
            qtensor: TchTensor::new(qtensor),
            scheme: *scheme,
        }
    }
    fn dequantize(tensor: QuantizedTensor<Self>) -> FloatTensor<Self> {
        TchTensor::new(tensor.qtensor.tensor.dequantize().to_kind(E::KIND))
    }
    fn q_device(tensor: &QuantizedTensor<Self>) -> LibTorchDevice {
        tensor.qtensor.tensor.device().into()
    }
    fn q_to_device(
        tensor: QuantizedTensor<Self>,
        device: &burn_tensor::Device<Self>,
    ) -> QuantizedTensor<Self> {
        let mut tensor = tensor;
        tensor.qtensor = TchOps::to_device(tensor.qtensor, device);
        tensor
    }
    fn q_reshape(tensor: QuantizedTensor<Self>, shape: Shape) -> QuantizedTensor<Self> {
        TchQTensor {
            qtensor: TchOps::reshape(tensor.qtensor, shape),
            scheme: tensor.scheme,
        }
    }
    async fn q_into_data(tensor: QuantizedTensor<Self>) -> TensorData {
        let shape = tensor.shape();
        let tensor = Self::q_reshape(tensor.clone(), Shape::new([shape.num_elements()]));
        let strategy = tensor.strategy();
        // To get the integer values we have to call `int_repr()`
        let values: Result<Vec<i8>, tch::TchError> = tensor.qtensor.tensor.int_repr().try_into();
        TensorData::quantized(values.unwrap(), shape, strategy)
    }
    fn q_swap_dims(
        tensor: QuantizedTensor<Self>,
        dim1: usize,
        dim2: usize,
    ) -> QuantizedTensor<Self> {
        // NOTE: with per-channel quantization (future), the channel axis could be impacted by this op
        let mut tensor = tensor;
        tensor.qtensor = TchOps::swap_dims(tensor.qtensor, dim1, dim2);
        tensor
    }
    fn q_permute(tensor: QuantizedTensor<Self>, axes: &[usize]) -> QuantizedTensor<Self> {
        // NOTE: with per-channel quantization (future), the channel axis could be impacted by this op
        let mut tensor = tensor;
        tensor.qtensor = TchOps::permute(tensor.qtensor, axes);
        tensor
    }
    fn q_flip(tensor: QuantizedTensor<Self>, axes: &[usize]) -> QuantizedTensor<Self> {
        let mut tensor = tensor;
        tensor.qtensor = TchOps::flip(tensor.qtensor, axes);
        tensor
    }
    fn q_select(
        tensor: QuantizedTensor<Self>,
        dim: usize,
        indices: IntTensor<Self>,
    ) -> QuantizedTensor<Self> {
        let mut tensor = tensor;
        tensor.qtensor = TchOps::index_select_dim(tensor.qtensor, dim, indices);
        tensor
    }
    fn q_slice(tensor: QuantizedTensor<Self>, ranges: &[Range<usize>]) -> QuantizedTensor<Self> {
        let mut tensor = tensor;
        tensor.qtensor = TchOps::slice(tensor.qtensor, ranges);
        tensor
    }
    fn q_argmax(tensor: QuantizedTensor<Self>, dim: usize) -> IntTensor<Self> {
        TchOps::argmax(TchTensor::new(tensor.qtensor.tensor.int_repr()), dim)
    }
    fn q_argmin(tensor: QuantizedTensor<Self>, dim: usize) -> IntTensor<Self> {
        TchOps::argmin(TchTensor::new(tensor.qtensor.tensor.int_repr()), dim)
    }
    fn q_max_dim_with_indices(
        tensor: QuantizedTensor<Self>,
        dim: usize,
    ) -> (QuantizedTensor<Self>, IntTensor<Self>) {
        let (qtensor, indices) = TchOps::max_dim_with_indices(tensor.qtensor, dim);
        let values = TchQTensor {
            qtensor,
            scheme: tensor.scheme,
        };
        (values, indices)
    }
    fn q_max_dim(tensor: QuantizedTensor<Self>, dim: usize) -> QuantizedTensor<Self> {
        TchQTensor {
            qtensor: TchOps::max_dim(tensor.qtensor, dim),
            scheme: tensor.scheme,
        }
    }
    fn q_min_dim(tensor: QuantizedTensor<Self>, dim: usize) -> QuantizedTensor<Self> {
        TchQTensor {
            qtensor: TchOps::min_dim(tensor.qtensor, dim),
            scheme: tensor.scheme,
        }
    }
    fn q_min_dim_with_indices(
        tensor: QuantizedTensor<Self>,
        dim: usize,
    ) -> (QuantizedTensor<Self>, IntTensor<Self>) {
        let (qtensor, indices) = TchOps::min_dim_with_indices(tensor.qtensor, dim);
        let values = TchQTensor {
            qtensor,
            scheme: tensor.scheme,
        };
        (values, indices)
    }
    fn q_expand(tensor: QuantizedTensor<Self>, shape: Shape) -> QuantizedTensor<Self> {
        // NOTE: with per-channel quantization (future), the channel axis could be impacted by this op
        TchQTensor {
            qtensor: TchOps::expand(tensor.qtensor, shape),
            scheme: tensor.scheme,
        }
    }
    fn q_sort(
        tensor: QuantizedTensor<Self>,
        dim: usize,
        descending: bool,
    ) -> QuantizedTensor<Self> {
        TchQTensor {
            qtensor: TchOps::sort(tensor.qtensor, dim, descending),
            scheme: tensor.scheme,
        }
    }
    fn q_sort_with_indices(
        tensor: QuantizedTensor<Self>,
        dim: usize,
        descending: bool,
    ) -> (QuantizedTensor<Self>, IntTensor<Self>) {
        let (qtensor, indices) = TchOps::sort_with_indices(tensor.qtensor, dim, descending);
        let tensor = TchQTensor {
            qtensor,
            scheme: tensor.scheme,
        };
        (tensor, indices)
    }
    fn q_argsort(tensor: QuantizedTensor<Self>, dim: usize, descending: bool) -> IntTensor<Self> {
        TchOps::argsort(tensor.qtensor, dim, descending)
    }
}
impl<E: TchElement, Q: QuantElement> FloatTensorOps<Self> for LibTorch<E, Q> {
    fn float_from_data(data: TensorData, device: &LibTorchDevice) -> TchTensor {
    fn float_random(
        shape: Shape,
        distribution: Distribution,
        device: &LibTorchDevice,
    ) -> TchTensor {
    fn float_repeat_dim(tensor: TchTensor, dim: usize, times: usize) -> TchTensor {
    fn float_zeros(shape: Shape, device: &LibTorchDevice) -> TchTensor {
    fn float_ones(shape: Shape, device: &LibTorchDevice) -> TchTensor {
    async fn float_into_data(tensor: TchTensor) -> TensorData {
    fn float_device(tensor: &TchTensor) -> LibTorchDevice {
    fn float_to_device(tensor: TchTensor, device: &LibTorchDevice) -> TchTensor {
    fn float_empty(shape: Shape, device: &<LibTorch<E> as Backend>::Device) -> TchTensor {
    fn float_add(lhs: TchTensor, rhs: TchTensor) -> TchTensor {
    fn float_add_scalar(lhs: TchTensor, rhs: E) -> TchTensor {
    fn float_sub(lhs: TchTensor, rhs: TchTensor) -> TchTensor {
    fn float_sub_scalar(lhs: TchTensor, rhs: E) -> TchTensor {
    fn float_mul(lhs: TchTensor, rhs: TchTensor) -> TchTensor {
    fn float_mul_scalar(lhs: TchTensor, rhs: E) -> TchTensor {
    fn float_div(lhs: TchTensor, rhs: TchTensor) -> TchTensor {
    fn float_div_scalar(lhs: TchTensor, rhs: E) -> TchTensor {
    fn float_remainder(lhs: TchTensor, rhs: TchTensor) -> TchTensor {
    fn float_remainder_scalar(lhs: TchTensor, rhs: E) -> TchTensor {
    fn float_matmul(lhs: TchTensor, rhs: TchTensor) -> TchTensor {
    fn float_neg(tensor: TchTensor) -> TchTensor {
    fn float_recip(tensor: TchTensor) -> TchTensor {
    fn float_swap_dims(tensor: TchTensor, dim1: usize, dim2: usize) -> TchTensor {
    fn float_reshape(tensor: TchTensor, shape: Shape) -> TchTensor {
    fn float_gather(dim: usize, tensor: TchTensor, indices: TchTensor) -> TchTensor {
    fn float_scatter(
        dim: usize,
        tensor: TchTensor,
        indices: TchTensor,
        value: TchTensor,
    ) -> TchTensor {
    fn float_select(tensor: TchTensor, dim: usize, indices: TchTensor) -> TchTensor {
    fn float_select_assign(
        tensor: TchTensor,
        dim: usize,
        indices: TchTensor,
        value: TchTensor,
    ) -> TchTensor {
    fn float_slice(tensor: TchTensor, ranges: &[Range<usize>]) -> TchTensor {
    fn float_slice_assign(
        tensor: TchTensor,
        ranges: &[Range<usize>],
        value: TchTensor,
    ) -> TchTensor {
    fn float_mask_where(tensor: TchTensor, mask: TchTensor, value: TchTensor) -> TchTensor {
    fn float_mask_fill(tensor: TchTensor, mask: TchTensor, value: E) -> TchTensor {
    fn float_equal(lhs: TchTensor, rhs: TchTensor) -> TchTensor {
    fn float_equal_elem(lhs: TchTensor, rhs: E) -> TchTensor {
    fn float_greater(lhs: TchTensor, rhs: TchTensor) -> TchTensor {
    fn float_greater_elem(lhs: TchTensor, rhs: E) -> TchTensor {
    fn float_greater_equal(lhs: TchTensor, rhs: TchTensor) -> TchTensor {
    fn float_greater_equal_elem(lhs: TchTensor, rhs: E) -> TchTensor {
    fn float_lower(lhs: TchTensor, rhs: TchTensor) -> TchTensor {
    fn float_lower_elem(lhs: TchTensor, rhs: E) -> TchTensor {
    fn float_lower_equal(lhs: TchTensor, rhs: TchTensor) -> TchTensor {
    fn float_lower_equal_elem(lhs: TchTensor, rhs: E) -> TchTensor {
    fn float_mean(tensor: TchTensor) -> TchTensor {
    fn float_sum(tensor: TchTensor) -> TchTensor {
    fn float_sum_dim(tensor: TchTensor, dim: usize) -> TchTensor {
    fn float_mean_dim(tensor: TchTensor, dim: usize) -> TchTensor {
    fn float_prod(tensor: TchTensor) -> TchTensor {
    fn float_prod_dim(tensor: TchTensor, dim: usize) -> TchTensor {
    fn float_argmax(tensor: TchTensor, dim: usize) -> TchTensor {
    fn float_argmin(tensor: TchTensor, dim: usize) -> TchTensor {
    fn float_max_dim(tensor: TchTensor, dim: usize) -> TchTensor {
    fn float_max_dim_with_indices(tensor: TchTensor, dim: usize) -> (TchTensor, TchTensor) {
    fn float_min_dim(tensor: TchTensor, dim: usize) -> TchTensor {
    fn float_min_dim_with_indices(tensor: TchTensor, dim: usize) -> (TchTensor, TchTensor) {
    fn float_exp(tensor: TchTensor) -> TchTensor {
    fn float_log(tensor: TchTensor) -> TchTensor {
    fn float_log1p(tensor: TchTensor) -> TchTensor {
    fn float_powf_scalar(tensor: TchTensor, value: f32) -> TchTensor {
    fn float_sqrt(tensor: TchTensor) -> TchTensor {
    fn float_abs(tensor: TchTensor) -> TchTensor {
    fn float_cos(tensor: TchTensor) -> TchTensor {
    fn float_sin(tensor: TchTensor) -> TchTensor {
    fn float_tanh(tensor: TchTensor) -> TchTensor {
    fn float_round(tensor: TchTensor) -> TchTensor {
    fn float_floor(tensor: TchTensor) -> TchTensor {
    fn float_ceil(tensor: TchTensor) -> TchTensor {
    fn float_erf(tensor: TchTensor) -> TchTensor {
    fn float_cat(tensors: Vec<TchTensor>, dim: usize) -> TchTensor {
    fn float_clamp_min(tensor: TchTensor, min: E) -> TchTensor {
    fn float_clamp_max(tensor: TchTensor, max: <LibTorch<E> as Backend>::FloatElem) -> TchTensor {
    fn float_clamp(
        tensor: TchTensor,
        min: <LibTorch<E> as Backend>::FloatElem,
        max: <LibTorch<E> as Backend>::FloatElem,
    ) -> TchTensor {
    fn float_into_int(tensor: TchTensor) -> TchTensor {
    fn float_powf(lhs: TchTensor, rhs: TchTensor) -> TchTensor {
    fn float_permute(tensor: TchTensor, axes: &[usize]) -> TchTensor {
    fn float_flip(tensor: TchTensor, axes: &[usize]) -> TchTensor {
    fn float_sign(tensor: TchTensor) -> TchTensor {
    fn float_expand(tensor: TchTensor, shape: Shape) -> TchTensor {
    fn float_sort(tensor: TchTensor, dim: usize, descending: bool) -> TchTensor {
    fn float_sort_with_indices(
        tensor: TchTensor,
        dim: usize,
        descending: bool,
    ) -> (TchTensor, TchTensor) {
    fn float_argsort(tensor: TchTensor, dim: usize, descending: bool) -> IntTensor<Self> {
    fn float_cast(tensor: TchTensor, dtype: FloatDType) -> TchTensor {
// NOTE: when dtypes of inputs to an arithmetic operation differ, tch handles type
// promotion based on a set of rules: https://pytorch.org/docs/stable/tensor_attributes.html#type-promotion-doc
// Type promotion is not automatic on all backends so this behavior might differ
impl<E: TchElement, Q: QuantElement> TransactionOps<Self> for LibTorch<E, Q> {
/// A reference to a tensor storage.
///
/// We manually implement `Sync` and `Send` unsafely, so even if we could use `Rc`, it isn't safe.
pub type StorageRef = Arc<*mut c_void>;
/// A reference to a tensor storage.
pub enum Storage {
/// When a tensor is a partial view of another tensor.
/// Storage reference for the whole buffer.
/// Storage reference for the partial buffer.
/// When a tensor use all of its buffer.
/// Storage reference for the whole buffer.
impl Storage {
/// Check if the storage can be used inplace.
    pub fn can_mut(&self) -> bool {
/// Get the whole buffer reference.
    pub fn buffer_ref(&self) -> &StorageRef {
/// A tensor using the tch backend.
pub struct TchTensor {
/// Handle to the tensor. Call methods on this field.
/// The tensor's storage
impl TensorMetadata for TchTensor {
    fn dtype(&self) -> DType {
// Complex and quantization types are not valid/implemented.
    fn shape(&self) -> Shape {
impl TchTensor {
/// Create a new tensor.
///
/// Note that if the tensor was created from an operation that may reuse the same tensor
/// storage as the parent, you should use [from_existing](TchTensor::from_existing)
/// instead.
    pub fn new(tensor: tch::Tensor) -> Self {
/// Create a tensor that was created from an operation executed on a parent tensor.
///
/// If the child tensor shared the same storage as its parent, it will be cloned, effectively
/// tracking how much tensors point to the same memory space.
    pub fn from_existing(tensor: tch::Tensor, storage_parent: Storage) -> Self {
/// Create a tensor that uses a part of its parent tensor such as slice and narrow.
    pub fn partial(tensor: tch::Tensor, storage_parent: Storage) -> Self {
// This is safe since we don't use autodiff from LibTorch.
// Also, atomic reference counting is used to know if the tensor's data can be reused.
// If there are multiple reference on the same tensor, it becomes read only.
unsafe impl Send for TchTensor {
unsafe impl Sync for TchTensor {
impl TchTensor {
/// Checks if the tensor can be mutated in-place.
///
/// Returns `true` if the tensor's stride does not contain zero (no broadcasting)
/// and the storage can be mutated.
    pub fn can_mut(&self) -> bool {
/// Executes an operation on a tensor if the data can be reused.
    pub fn mut_ops<F: Fn(&mut tch::Tensor) -> tch::Tensor>(
        &mut self,
        func: F,
    ) -> Option<TchTensor> {
/// Executes a unary operation, reusing the tensor data if possible.
    pub fn unary_ops<FOwn, FRef>(self, fown: FOwn, fref: FRef) -> TchTensor
    where
        FOwn: Fn(tch::Tensor) -> tch::Tensor,
        FRef: Fn(&tch::Tensor) -> tch::Tensor,
    {
/// Executes a binary operation, reusing the tensor data if possible.
    pub fn binary_ops_tensor<FLMut, FRMut, FRef>(
        mut lhs: Self,
        mut rhs: Self,
        flmut: FLMut,
        frmut: FRMut,
        fref: FRef,
    ) -> TchTensor
    where
        FLMut: Fn(&mut tch::Tensor, &tch::Tensor) -> tch::Tensor,
        FRMut: Fn(&tch::Tensor, &mut tch::Tensor) -> tch::Tensor,
        FRef: Fn(&tch::Tensor, &tch::Tensor) -> tch::Tensor,
    {
// Both lhs and rhs are expected to have the same rank
// Attempt to mutate lhs tensor
// Attempt to mutate rhs tensor
impl Clone for TchTensor {
    fn clone(&self) -> Self {
/// A shape that can be used by LibTorch.
pub struct TchShape {
/// The shape's dimensions.
impl From<Shape> for TchShape {
    fn from(shape: Shape) -> Self {
impl From<&[usize]> for TchShape {
    fn from(shape: &[usize]) -> Self {
impl TchTensor {
/// Creates a new tensor from a shape and a device.
///
/// # Arguments
///
/// * `data` - The tensor's data.
/// * `device` - The device on which the tensor will be allocated.
///
/// # Returns
///
/// A new tensor.
    pub fn from_data<E: TchElement>(data: TensorData, device: tch::Device) -> Self {
impl TchTensor {
/// Creates an empty tensor from a shape and a device.
///
/// # Arguments
///
/// * `shape` - The shape of the tensor.
/// * `device` - The device to create the tensor on.
///
/// # Returns
///
/// A new empty tensor.
    pub fn empty<E: tch::kind::Element>(shape: Shape, device: LibTorchDevice) -> Self {
/// A quantized tensor for the tch backend.
pub struct TchQTensor {
/// The quantized tensor.
/// The quantization scheme.
impl TchQTensor {
/// Returns the quantization strategy, including quantization parameters, for the given tensor.
    pub fn strategy(&self) -> QuantizationStrategy {
impl TensorMetadata for TchQTensor {
    fn dtype(&self) -> DType {
    fn shape(&self) -> Shape {
impl QTensorPrimitive for TchQTensor {
    fn scheme(&self) -> &QuantScheme {
    fn should_support_into_and_from_data_1d() {
    fn should_support_into_and_from_data_2d() {
    fn should_not_update_inplace_after_reshape() {
    fn should_not_update_inplace_after_slice() {
    fn should_support_qtensor_strategy() {
/// Dummy function to get CUDA to link properly
    pub fn dummy_cuda_dependency();
static INIT_ARRAY: [unsafe extern "C" fn();
enum Message<R, B: Backend> {
struct CheckpointerThread<C, R, B: Backend> {
impl<C, R, B> CheckpointerThread<C, R, B>
where
    C: Checkpointer<R, B>,
    R: Record<B>,
    B: Backend,
{
    fn run(self) {
/// Async checkpointer.
pub struct AsyncCheckpointer<Record, B: Backend> {
impl<R, B> AsyncCheckpointer<R, B>
where
    R: Record<B> + 'static,
    B: Backend,
{
/// Create a new async checkpointer.
///
/// # Arguments
///
/// * `checkpointer` - The checkpointer.
///
/// # Returns
///
/// The async checkpointer.
    pub fn new<C>(checkpointer: C) -> Self
    where
        C: Checkpointer<R, B> + Send + 'static,
    {
// Only on checkpoint can be done in advance.
impl<R, B> Checkpointer<R, B> for AsyncCheckpointer<R, B>
where
    R: Record<B> + 'static,
    B: Backend,
{
    fn save(&self, epoch: usize, record: R) -> Result<(), CheckpointerError> {
    fn restore(&self, epoch: usize, device: &B::Device) -> Result<R, CheckpointerError> {
    fn delete(&self, epoch: usize) -> Result<(), CheckpointerError> {
impl<E, B> Drop for AsyncCheckpointer<E, B>
where
    B: Backend,
{
    fn drop(&mut self) {
/// The error type for checkpointer.
pub enum CheckpointerError {
/// IO error.
/// Recorder error.
/// Other errors.
/// The trait for checkpointer.
pub trait Checkpointer<R, B>
where
    R: Record<B>,
    B: Backend,
{
/// Save the record.
///
/// # Arguments
///
/// * `epoch` - The epoch.
/// * `record` - The record.
    fn save(&self, epoch: usize, record: R) -> Result<(), CheckpointerError>;
/// Delete the record at the given epoch if present.
    fn delete(&self, epoch: usize) -> Result<(), CheckpointerError>;
/// Restore the record.
///
/// # Arguments
///
/// * `epoch` - The epoch.
/// * `device` - The device used to restore the record.
///
/// # Returns
///
/// The record.
    fn restore(&self, epoch: usize, device: &B::Device) -> Result<R, CheckpointerError>;
/// The file checkpointer.
pub struct FileCheckpointer<FR> {
impl<FR> FileCheckpointer<FR> {
/// Creates a new file checkpointer.
///
/// # Arguments
///
/// * `recorder` - The file recorder.
/// * `directory` - The directory to save the checkpoints.
/// * `name` - The name of the checkpoint.
    pub fn new(recorder: FR, directory: impl AsRef<Path>, name: &str) -> Self {
    fn path_for_epoch(&self, epoch: usize) -> PathBuf {
impl<FR, R, B> Checkpointer<R, B> for FileCheckpointer<FR>
where
    R: Record<B>,
    FR: FileRecorder<B>,
    B: Backend,
{
    fn save(&self, epoch: usize, record: R) -> Result<(), CheckpointerError> {
    fn restore(&self, epoch: usize, device: &B::Device) -> Result<R, CheckpointerError> {
    fn delete(&self, epoch: usize) -> Result<(), CheckpointerError> {
/// Action to be taken by a [checkpointer](crate::checkpoint::Checkpointer).
pub enum CheckpointingAction {
/// Delete the given epoch.
/// Save the current record.
/// Define when checkpoint should be saved and deleted.
pub trait CheckpointingStrategy {
/// Based on the epoch, determine if the checkpoint should be saved.
    fn checkpointing(
        &mut self,
        epoch: usize,
        collector: &EventStoreClient,
    ) -> Vec<CheckpointingAction>;
// We make dyn box implement the checkpointing strategy so that it can be used with generic, but
// still be dynamic.
impl CheckpointingStrategy for Box<dyn CheckpointingStrategy> {
    fn checkpointing(
        &mut self,
        epoch: usize,
        collector: &EventStoreClient,
    ) -> Vec<CheckpointingAction> {
/// Compose multiple checkpointing strategy and only delete checkpoints when both strategy flag an
/// epoch to be deleted.
pub struct ComposedCheckpointingStrategy {
/// Help building a [checkpointing strategy](CheckpointingStrategy) by combining multiple ones.
pub struct ComposedCheckpointingStrategyBuilder {
impl ComposedCheckpointingStrategyBuilder {
/// Add a new [checkpointing strategy](CheckpointingStrategy).
    pub fn add<S>(mut self, strategy: S) -> Self
    where
        S: CheckpointingStrategy + 'static,
    {
/// Create a new [composed checkpointing strategy](ComposedCheckpointingStrategy).
    pub fn build(self) -> ComposedCheckpointingStrategy {
impl ComposedCheckpointingStrategy {
    fn new(strategies: Vec<Box<dyn CheckpointingStrategy>>) -> Self {
/// Create a new builder which help compose multiple
/// [checkpointing strategies](CheckpointingStrategy).
    pub fn builder() -> ComposedCheckpointingStrategyBuilder {
impl CheckpointingStrategy for ComposedCheckpointingStrategy {
    fn checkpointing(
        &mut self,
        epoch: usize,
        collector: &EventStoreClient,
    ) -> Vec<CheckpointingAction> {
// We assume that the strategy would not want the current epoch to be saved.
// So we flag it as deleted.
    fn should_delete_when_both_deletes() {
/// Keep the last N checkpoints.
///
/// Very useful when training, minimizing disk space while ensuring that the training can be
/// resumed even if something goes wrong.
pub struct KeepLastNCheckpoints {
impl CheckpointingStrategy for KeepLastNCheckpoints {
    fn checkpointing(
        &mut self,
        epoch: usize,
        _store: &EventStoreClient,
    ) -> Vec<CheckpointingAction> {
    fn should_always_delete_lastn_epoch_if_higher_than_one() {
/// Keep the best checkpoint based on a metric.
pub struct MetricCheckpointingStrategy {
impl MetricCheckpointingStrategy {
/// Create a new metric checkpointing strategy.
    pub fn new<M>(metric: &M, aggregate: Aggregate, direction: Direction, split: Split) -> Self
    where
        M: Metric,
    {
impl CheckpointingStrategy for MetricCheckpointingStrategy {
    fn checkpointing(
        &mut self,
        epoch: usize,
        store: &EventStoreClient,
    ) -> Vec<CheckpointingAction> {
    fn always_keep_the_best_epoch() {
// Register an in memory logger.
// Register the loss metric.
// Two points for the first epoch. Mean 0.75
// Should save the current record.
// Two points for the second epoch. Mean 0.4
// Should save the current record and delete the previous one.
// Two points for the last epoch. Mean 2.0
// Should not delete the previous record, since it's the best one, and should not save a
// new one.
/// All components necessary to train a model grouped in one trait.
pub trait LearnerComponents {
/// The backend in used for the training.
    type Backend: AutodiffBackend;
/// The learning rate scheduler used for the training.
    type LrScheduler: LrScheduler;
/// The model to train.
    type Model: AutodiffModule<Self::Backend> + core::fmt::Display + 'static;
/// The optimizer used for the training.
    type Optimizer: Optimizer<Self::Model, Self::Backend>;
/// The checkpointer used for the model.
    type CheckpointerModel: Checkpointer<<Self::Model as Module<Self::Backend>>::Record, Self::Backend>;
/// The checkpointer used for the optimizer.
    type CheckpointerOptimizer: Checkpointer<
            <Self::Optimizer as Optimizer<Self::Model, Self::Backend>>::Record,
            Self::Backend,
        >;
/// The checkpointer used for the scheduler.
    type CheckpointerLrScheduler: Checkpointer<<Self::LrScheduler as LrScheduler>::Record<Self::Backend>, Self::Backend>;
    type EventProcessor: EventProcessor + 'static;
/// The strategy to save and delete checkpoints.
    type CheckpointerStrategy: CheckpointingStrategy;
/// Concrete type that implements [training components trait](TrainingComponents).
pub struct LearnerComponentsMarker<B, LR, M, O, CM, CO, CS, EP, S> {
impl<B, LR, M, O, CM, CO, CS, EP, S> LearnerComponents
    for LearnerComponentsMarker<B, LR, M, O, CM, CO, CS, EP, S>
where
    B: AutodiffBackend,
    LR: LrScheduler,
    M: AutodiffModule<B> + core::fmt::Display + 'static,
    O: Optimizer<M, B>,
    CM: Checkpointer<M::Record, B>,
    CO: Checkpointer<O::Record, B>,
    CS: Checkpointer<LR::Record<B>, B>,
    EP: EventProcessor + 'static,
    S: CheckpointingStrategy,
{
    type Backend = B;
    type LrScheduler = LR;
    type Model = M;
    type Optimizer = O;
    type CheckpointerModel = CM;
    type CheckpointerOptimizer = CO;
    type CheckpointerLrScheduler = CS;
    type EventProcessor = EP;
    type CheckpointerStrategy = S;
/// The training backend.
pub type TrainBackend<LC> = <LC as LearnerComponents>::Backend;
/// The validation backend.
pub type ValidBackend<LC> = <<LC as LearnerComponents>::Backend as AutodiffBackend>::InnerBackend;
/// This trait is used to install an application logger.
pub trait ApplicationLoggerInstaller {
/// Install the application logger.
    fn install(&self) -> Result<(), String>;
/// This struct is used to install a local file application logger to output logs to a given file path.
pub struct FileApplicationLoggerInstaller {
impl FileApplicationLoggerInstaller {
/// Create a new file application logger.
    pub fn new(path: impl AsRef<Path>) -> Self {
impl ApplicationLoggerInstaller for FileApplicationLoggerInstaller {
    fn install(&self) -> Result<(), String> {
// The wgpu crate is logging too much, so we skip `info` level.
/// Learner struct encapsulating all components necessary to train a Neural Network model.
///
/// To create a learner, use the [builder](crate::learner::LearnerBuilder) struct.
pub struct Learner<LC: LearnerComponents> {
pub(crate) struct LearnerCheckpointer<LC: LearnerComponents> {
impl<LC: LearnerComponents> LearnerCheckpointer<LC> {
    pub(crate) fn checkpoint(
        &mut self,
        model: &LC::Model,
        optim: &LC::Optimizer,
        scheduler: &LC::LrScheduler,
        epoch: usize,
        store: &EventStoreClient,
    ) {
    pub(crate) fn load_checkpoint(
        &self,
        model: LC::Model,
        optim: LC::Optimizer,
        scheduler: LC::LrScheduler,
        device: &Device<LC::Backend>,
        epoch: usize,
    ) -> (LC::Model, LC::Optimizer, LC::LrScheduler) {
/// A handle that allows aborting the training process early.
pub struct TrainingInterrupter {
impl TrainingInterrupter {
/// Create a new instance.
    pub fn new() -> Self {
/// Notify the learner that it should stop.
    pub fn stop(&self) {
/// True if .stop() has been called.
    pub fn should_stop(&self) -> bool {
/// Struct to configure and create a [learner](Learner).
pub struct LearnerBuilder<B, T, V, M, O, S>
where
    T: ItemLazy + 'static,
    V: ItemLazy + 'static,
    B: AutodiffBackend,
    M: AutodiffModule<B>,
    O: Optimizer<M, B>,
    S: LrScheduler,
{
// Not that complex and very convenient when the traits are
// already constrained correctly. Extracting in another type
// would be more complex.
    renderer: Option<Box<dyn MetricsRenderer + 'static>>,
    metrics: Metrics<T, V>,
    event_store: LogEventStore,
    interrupter: TrainingInterrupter,
    tracing_logger: Option<Box<dyn ApplicationLoggerInstaller>>,
    num_loggers: usize,
    checkpointer_strategy: Box<dyn CheckpointingStrategy>,
    early_stopping: Option<Box<dyn EarlyStoppingStrategy>>,
    // Use BTreeSet instead of HashSet for consistent (alphabetical) iteration order
    summary_metrics: BTreeSet<String>,
    summary: bool,
}

impl<B, T, V, M, O, S> LearnerBuilder<B, T, V, M, O, S>
where
    B: AutodiffBackend,
    T: ItemLazy + 'static,
    V: ItemLazy + 'static,
    M: AutodiffModule<B> + core::fmt::Display + 'static,
    O: Optimizer<M, B>,
    S: LrScheduler,
{
/// Creates a new learner builder.
///
/// # Arguments
///
/// * `directory` - The directory to save the checkpoints.
    pub fn new(directory: impl AsRef<Path>) -> Self {
// default to valid loss
/// Replace the default metric loggers with the provided ones.
///
/// # Arguments
///
/// * `logger_train` - The training logger.
/// * `logger_valid` - The validation logger.
    pub fn metric_loggers<MT, MV>(mut self, logger_train: MT, logger_valid: MV) -> Self
    where
        MT: MetricLogger + 'static,
        MV: MetricLogger + 'static,
    {
/// Update the checkpointing_strategy.
    pub fn with_checkpointing_strategy<CS>(mut self, strategy: CS) -> Self
    where
        CS: CheckpointingStrategy + 'static,
    {
/// Replace the default CLI renderer with a custom one.
///
/// # Arguments
///
/// * `renderer` - The custom renderer.
    pub fn renderer<MR>(mut self, renderer: MR) -> Self
    where
        MR: MetricsRenderer + 'static,
    {
/// Register a training metric.
    pub fn metric_train<Me: Metric + 'static>(mut self, metric: Me) -> Self
    where
        T::ItemSync: Adaptor<Me::Input>,
    {
/// Register a validation metric.
    pub fn metric_valid<Me: Metric + 'static>(mut self, metric: Me) -> Self
    where
        V::ItemSync: Adaptor<Me::Input>,
    {
/// Enable gradients accumulation.
///
/// # Notes
///
/// When you enable gradients accumulation, the gradients object used by the optimizer will be
/// the sum of all gradients generated by each backward pass. It might be a good idea to
/// reduce the learning to compensate.
///
/// The effect is similar to increasing the `batch size` and the `learning rate` by the `accumulation`
/// amount.
    pub fn grads_accumulation(mut self, accumulation: usize) -> Self {
/// Register a [numeric](crate::metric::Numeric) training [metric](Metric).
    pub fn metric_train_numeric<Me>(mut self, metric: Me) -> Self
    where
        Me: Metric + crate::metric::Numeric + 'static,
        T::ItemSync: Adaptor<Me::Input>,
    {
/// Register a [numeric](crate::metric::Numeric) validation [metric](Metric).
    pub fn metric_valid_numeric<Me: Metric + crate::metric::Numeric + 'static>(
        mut self,
        metric: Me,
    ) -> Self
    where
        V::ItemSync: Adaptor<Me::Input>,
    {
/// The number of epochs the training should last.
    pub fn num_epochs(mut self, num_epochs: usize) -> Self {
/// Run the training loop on multiple devices.
    pub fn devices(mut self, devices: Vec<B::Device>) -> Self {
/// The epoch from which the training must resume.
    pub fn checkpoint(mut self, checkpoint: usize) -> Self {
/// Provides a handle that can be used to interrupt training.
    pub fn interrupter(&self) -> TrainingInterrupter {
/// Register an [early stopping strategy](EarlyStoppingStrategy) to stop the training when the
/// conditions are meet.
    pub fn early_stopping<Strategy>(mut self, strategy: Strategy) -> Self
    where
        Strategy: EarlyStoppingStrategy + 'static,
    {
/// By default, Rust logs are captured and written into
/// `experiment.log`. If disabled, standard Rust log handling
/// will apply.
    pub fn with_application_logger(
        mut self,
        logger: Option<Box<dyn ApplicationLoggerInstaller>>,
    ) -> Self {
/// Register a checkpointer that will save the [optimizer](Optimizer), the
/// [model](AutodiffModule) and the [scheduler](LrScheduler) to different files.
    pub fn with_file_checkpointer<FR>(mut self, recorder: FR) -> Self
    where
        FR: FileRecorder<B> + 'static,
        FR: FileRecorder<B::InnerBackend> + 'static,
        O::Record: 'static,
        M::Record: 'static,
        S::Record<B>: 'static,
    {
/// Enable the training summary report.
///
/// The summary will be displayed at the end of `.fit()`.
    pub fn summary(mut self) -> Self {
/// Create the [learner](Learner) from a [model](AutodiffModule) and an [optimizer](Optimizer).
/// The [learning rate scheduler](LrScheduler) can also be a simple
/// [learning rate](burn_core::LearningRate).
// The goal for the builder is to handle all types and
// creates a clean learner.
    pub fn build(
        mut self,
        model: M,
        optim: O,
        lr_scheduler: S,
    ) -> Learner<
        LearnerComponentsMarker<
            B,
            S,
            M,
            O,
            AsyncCheckpointer<M::Record, B>,
            AsyncCheckpointer<O::Record, B>,
            AsyncCheckpointer<S::Record<B>, B>,
            AsyncProcessor<FullEventProcessor<T, V>>,
            Box<dyn CheckpointingStrategy>,
        >,
    >
    where
        M::Record: 'static,
        O::Record: 'static,
        S::Record<B>: 'static,
    {
/// Simple classification output adapted for multiple metrics.
pub struct ClassificationOutput<B: Backend> {
/// The loss.
/// The output.
/// The targets.
impl<B: Backend> ItemLazy for ClassificationOutput<B> {
    type ItemSync = ClassificationOutput<NdArray>;
    fn sync(self) -> Self::ItemSync {
impl<B: Backend> Adaptor<AccuracyInput<B>> for ClassificationOutput<B> {
    fn adapt(&self) -> AccuracyInput<B> {
impl<B: Backend> Adaptor<LossInput<B>> for ClassificationOutput<B> {
    fn adapt(&self) -> LossInput<B> {
impl<B: Backend> Adaptor<TopKAccuracyInput<B>> for ClassificationOutput<B> {
    fn adapt(&self) -> TopKAccuracyInput<B> {
impl<B: Backend> Adaptor<ConfusionStatsInput<B>> for ClassificationOutput<B> {
    fn adapt(&self) -> ConfusionStatsInput<B> {
/// Multi-label classification output adapted for multiple metrics.
pub struct MultiLabelClassificationOutput<B: Backend> {
/// The loss.
/// The output.
/// The targets.
impl<B: Backend> ItemLazy for MultiLabelClassificationOutput<B> {
    type ItemSync = MultiLabelClassificationOutput<NdArray>;
    fn sync(self) -> Self::ItemSync {
impl<B: Backend> Adaptor<HammingScoreInput<B>> for MultiLabelClassificationOutput<B> {
    fn adapt(&self) -> HammingScoreInput<B> {
impl<B: Backend> Adaptor<LossInput<B>> for MultiLabelClassificationOutput<B> {
    fn adapt(&self) -> LossInput<B> {
impl<B: Backend> Adaptor<ConfusionStatsInput<B>> for MultiLabelClassificationOutput<B> {
    fn adapt(&self) -> ConfusionStatsInput<B> {
/// The condition that [early stopping strategies](EarlyStoppingStrategy) should follow.
pub enum StoppingCondition {
/// When no improvement has happened since the given number of epochs.
/// The number of epochs allowed to worsen before it gets better.
/// A strategy that checks if the training should be stopped.
pub trait EarlyStoppingStrategy {
/// Update its current state and returns if the training should be stopped.
    fn should_stop(&mut self, epoch: usize, store: &EventStoreClient) -> bool;
/// An [early stopping strategy](EarlyStoppingStrategy) based on a metrics collected
/// during training or validation.
pub struct MetricEarlyStoppingStrategy {
impl EarlyStoppingStrategy for MetricEarlyStoppingStrategy {
    fn should_stop(&mut self, epoch: usize, store: &EventStoreClient) -> bool {
impl MetricEarlyStoppingStrategy {
/// Create a new [early stopping strategy](EarlyStoppingStrategy) based on a metrics collected
/// during training or validation.
///
/// # Notes
///
/// The metric should be registered for early stopping to work, otherwise no data is collected.
    pub fn new<Me: Metric>(
        metric: &Me,
        aggregate: Aggregate,
        direction: Direction,
        split: Split,
        condition: StoppingCondition,
    ) -> Self {
    fn never_early_stop_while_it_is_improving() {
    fn early_stop_when_no_improvement_since_two_epochs() {
    fn early_stop_when_stays_equal() {
    fn test_early_stopping(n_epochs: usize, data: &[(&[f64], bool, &str)]) {
/// A validation epoch.
pub struct ValidEpoch<B: Backend, VI> {
/// A training epoch.
pub struct TrainEpoch<B: AutodiffBackend, TI> {
impl<B: Backend, VI> ValidEpoch<B, VI> {
/// Runs the validation epoch.
///
/// # Arguments
///
/// * `model` - The model to validate.
/// * `processor` - The event processor to use.
    pub fn run<LC: LearnerComponents, VO>(
        &self,
        model: &LC::Model,
        processor: &mut LC::EventProcessor,
        interrupter: &TrainingInterrupter,
    ) where
        LC::EventProcessor: EventProcessor<ItemValid = VO>,
        <LC::Model as AutodiffModule<LC::Backend>>::InnerModule: ValidStep<VI, VO>,
        LC::Backend: AutodiffBackend<InnerBackend = B>,
    {
impl<B: AutodiffBackend, TI> TrainEpoch<B, TI> {
/// Runs the training epoch.
///
/// # Arguments
///
/// * `model` - The model to train.
/// * `optim` - The optimizer to use.
/// * `scheduler` - The learning rate scheduler to use.
/// * `processor` - The event processor to use.
///
/// # Returns
///
/// The trained model and the optimizer.
    pub fn run<LC: LearnerComponents<Backend = B>, TO>(
        &mut self,
        mut model: LC::Model,
        mut optim: LC::Optimizer,
        scheduler: &mut LC::LrScheduler,
        processor: &mut LC::EventProcessor,
        interrupter: &TrainingInterrupter,
    ) -> (LC::Model, LC::Optimizer)
    where
        LC::EventProcessor: EventProcessor<ItemTrain = TO>,
        LC::Model: TrainStep<TI, TO>,
    {
// Single device / dataloader
impl<B: AutodiffBackend, TI> TrainEpoch<B, TI> {
/// Runs the training epoch on multiple devices.
///
/// # Arguments
///
/// * `model` - The model to train.
/// * `optim` - The optimizer to use.
/// * `lr_scheduler` - The learning rate scheduler to use.
/// * `processor` - The event processor to use.
/// * `devices` - The devices to use.
///
/// # Returns
///
/// The trained model and the optimizer.
    pub fn run_multi_device<LC: LearnerComponents<Backend = B>, TO>(
        &mut self,
        mut model: LC::Model,
        mut optim: LC::Optimizer,
        lr_scheduler: &mut LC::LrScheduler,
        processor: &mut LC::EventProcessor,
        devices: Vec<<LC::Backend as Backend>::Device>,
        interrupter: &TrainingInterrupter,
    ) -> (LC::Model, LC::Optimizer)
    where
        LC::EventProcessor: EventProcessor<ItemTrain = TO>,
        LC::Model: TrainStep<TI, TO>,
        TO: Send + 'static,
        TI: Send + 'static,
    {
// The main device is always the first in the list.
// TODO: aggregate multi device (all-reduce)
/// Simple regression output adapted for multiple metrics.
pub struct RegressionOutput<B: Backend> {
/// The loss.
/// The output.
/// The targets.
impl<B: Backend> Adaptor<LossInput<B>> for RegressionOutput<B> {
    fn adapt(&self) -> LossInput<B> {
impl<B: Backend> ItemLazy for RegressionOutput<B> {
    type ItemSync = RegressionOutput<NdArray>;
    fn sync(self) -> Self::ItemSync {
/// The trainer module.
/// Multi devices train step.
pub struct MultiDevicesTrainStep<B: AutodiffBackend, M, TI, TO> {
struct Message<M, TI> {
struct Worker<B: AutodiffBackend, M, TI> {
impl<B, M, TI> Worker<B, M, TI>
where
    B: AutodiffBackend,
    M: AutodiffModule<B>,
{
    fn register(&self, item: TI, model: &M) {
    fn start<TO>(
        &self,
        sender_output: Sender<TrainOutput<TO>>,
        receiver_input: Receiver<Message<M, TI>>,
    ) where
        TI: Send + 'static,
        TO: Send + 'static,
        M: TrainStep<TI, TO> + Send + 'static,
    {
impl<B, M, TI, TO> MultiDevicesTrainStep<B, M, TI, TO>
where
    B: AutodiffBackend,
    M: AutodiffModule<B> + TrainStep<TI, TO> + Send + Clone + 'static,
    TI: Send + 'static,
    TO: Send + 'static,
{
/// Create a new multi devices train step.
///
/// # Arguments
///
/// * `devices` - Devices.
///
/// # Returns
///
/// MultiDevicesTrainStep instance.
    pub fn new(devices: &[B::Device]) -> Self
    where
        TI: Send + 'static,
    {
/// Collect outputs from workers for one step.
///
/// # Arguments
///
/// * `model` - Model.
/// * `dataloaders` - The data loader for each worker.
///
/// # Returns
///
/// Outputs.
    pub fn step<'a>(
        &self,
        dataloaders: &mut [Box<dyn DataLoaderIterator<TI> + 'a>],
        model: &M,
    ) -> (Vec<TrainOutput<TO>>, Progress) {
/// Contains the metric value at a given time.
pub struct MetricEntry {
/// The step at which the metric was recorded (i.e., epoch).
/// The metric value.
/// Contains the summary of recorded values for a given metric.
pub struct MetricSummary {
/// The metric name.
/// The metric entries.
impl MetricSummary {
    fn new<E: EventStore>(
        event_store: &mut E,
        metric: &str,
        split: Split,
        num_epochs: usize,
    ) -> Option<Self> {
/// Contains the summary of recorded metrics for the training and validation steps.
pub struct SummaryMetrics {
/// Training metrics summary.
/// Validation metrics summary.
/// Detailed training summary.
pub struct LearnerSummary {
/// The number of epochs completed.
/// The summary of recorded metrics during training.
/// The model name (only recorded within the learner).
impl LearnerSummary {
/// Creates a new learner summary for the specified metrics.
///
/// # Arguments
///
/// * `directory` - The directory containing the training artifacts (checkpoints and logs).
/// * `metrics` - The list of metrics to collect for the summary.
    pub fn new<S: AsRef<str>>(directory: impl AsRef<Path>, metrics: &[S]) -> Result<Self, String> {
// Number of recorded epochs
    pub(crate) fn with_model(mut self, name: String) -> Self {
impl Display for LearnerSummary {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
// Compute the max length for each column
// Summary header
// Metrics table header
// Table entries
        fn cmp_f64(a: &f64, b: &f64) -> Ordering {
        fn fmt_val(val: f64) -> String {
// Use scientific notation for small values which would otherwise be truncated
// skip metrics with no recorded values
// Compute the min & max for each metric
pub(crate) struct LearnerSummaryConfig {
impl LearnerSummaryConfig {
    pub fn init(&self) -> Result<LearnerSummary, String> {
    fn test_artifact_dir_should_exist() {
    fn test_train_valid_artifacts_should_exist() {
    fn test_summary_should_be_empty() {
    fn test_summary_should_be_collected() {
// Only Loss metric
// Aggregated train metric entries for 1 epoch
// epoch = 1
// (1 + 2) / 2
// Aggregated valid metric entries for 1 epoch
// epoch = 1
/// A training output.
pub struct TrainOutput<TO> {
/// The gradients.
/// The item.
impl<TO> TrainOutput<TO> {
/// Creates a new training output.
///
/// # Arguments
///
/// * `module` - The module.
/// * `grads` - The gradients.
/// * `item` - The item.
///
/// # Returns
///
/// A new training output.
    pub fn new<B: AutodiffBackend, M: AutodiffModule<B>>(
        module: &M,
        grads: B::Gradients,
        item: TO,
    ) -> Self {
/// Trait to be implemented for training models.
///
/// The [step](TrainStep::step) method needs to be manually implemented for all structs.
///
/// The [optimize](TrainStep::optimize) method can be overridden if you want to control how the
/// optimizer is used to update the model. This can be useful if you want to call custom mutable
/// functions on your model (e.g., clipping the weights) before or after the optimizer is used.
///
/// # Notes
///
/// To be used with the [Learner](Learner) struct, the struct which implements this trait must
/// also implement the [AutodiffModule] trait, which is done automatically with the
/// [Module](burn_core::module::Module) derive.
pub trait TrainStep<TI, TO> {
/// Runs the training step, which executes the forward and backward passes.
///
/// # Arguments
///
/// * `item` - The training input for the model.
///
/// # Returns
///
/// The training output containing the model output and the gradients.
    fn step(&self, item: TI) -> TrainOutput<TO>;
/// Optimize the current module with the provided gradients and learning rate.
///
/// # Arguments
///
/// * `optim`: Optimizer used for training this model.
/// * `lr`: The learning rate used for this step.
/// * `grads`: The gradients of each parameter in the current model.
///
/// # Returns
///
/// The updated model.
    fn optimize<B, O>(self, optim: &mut O, lr: f64, grads: GradientsParams) -> Self
    where
        B: AutodiffBackend,
        O: Optimizer<Self, B>,
        Self: AutodiffModule<B>,
    {
/// Trait to be implemented for validating models.
pub trait ValidStep<VI, VO> {
/// Runs a validation step.
///
/// # Arguments
///
/// * `item` - The item to validate on.
///
/// # Returns
///
/// The validation output.
    fn step(&self, item: VI) -> VO;
impl<LC: LearnerComponents> Learner<LC> {
/// Fits the model.
///
/// # Arguments
///
/// * `dataloader_train` - The training dataloader.
/// * `dataloader_valid` - The validation dataloader.
///
/// # Returns
///
/// The fitted model.
    pub fn fit<InputTrain, InputValid, OutputTrain, OutputValid>(
        mut self,
        mut dataloader_train: Arc<dyn DataLoader<TrainBackend<LC>, InputTrain>>,
        mut dataloader_valid: Arc<dyn DataLoader<ValidBackend<LC>, InputValid>>,
    ) -> LC::Model
    where
        InputTrain: Send + 'static,
        InputValid: Send,
        OutputTrain: Send + 'static,
        OutputValid: Send,
        LC::Model: TrainStep<InputTrain, OutputTrain>,
        <LC::Model as AutodiffModule<LC::Backend>>::InnerModule: ValidStep<InputValid, OutputValid>,
        LC::EventProcessor: EventProcessor<ItemTrain = OutputTrain, ItemValid = OutputValid>,
    {
// The reference model is always on the first device provided.
// Load the checkpoint on the default device.
// `MultiDevicesTrainStep` has one worker per device, so we use a fixed device strategy
// for each (worker) data loader. This matches the expected device on the worker, so we
// don't have to move the data between devices.
// Changed the train epoch to keep the dataloaders
// TODO: multi-device validation?
// Signal training end. For the TUI renderer, this handles the exit & return to main screen.
// Display learner summary
//! A library for training neural networks using the burn crate.
/// The checkpoint module.
/// Renderer modules to display metrics and training information.
/// The logger module.
/// The metric module.
pub(crate) type TestBackend = burn_ndarray::NdArray<f32>;
/// Probability of tp before adding errors
    pub const THRESHOLD: f64 = 0.5;
    pub enum ClassificationType {
/// Sample x Class shaped matrix for use in
/// classification metrics testing
    pub fn dummy_classification_input(
        classification_type: &ClassificationType,
    ) -> (Tensor<TestBackend, 2>, Tensor<TestBackend, 2, Bool>) {
// targets
// predictions @ threshold=0.5
//                     [[0], [0], [1], [0], [1]]
// targets
// predictions @ top_k=1
//   [[0, 1, 0], [0, 1, 0], [1, 0, 0], [0, 0, 1], [1, 0,  0]]
// predictions @ top_k=2
//   [[1, 1, 0], [1, 1, 0], [1, 1, 0], [0, 1, 1], [1, 0,  1]]
// targets
// predictions @ threshold=0.5
//   [[0, 1, 1], [0, 1, 0], [1, 1, 0], [1, 0, 1], [1, 0, 0]]
enum Message<T> {
/// Async logger.
pub struct AsyncLogger<T> {
struct LoggerThread<T, L: Logger<T>> {
impl<T, L> LoggerThread<T, L>
where
    L: Logger<T>,
{
    fn run(mut self) {
impl<T: Send + Sync + 'static> AsyncLogger<T> {
/// Create a new async logger.
    pub fn new<L>(logger: L) -> Self
    where
        L: Logger<T> + 'static,
    {
/// Sync the async logger.
    pub(crate) fn sync(&self) {
impl<T: Send> Logger<T> for AsyncLogger<T> {
    fn log(&mut self, item: T) {
impl<T> Drop for AsyncLogger<T> {
    fn drop(&mut self) {
/// The logger trait.
pub trait Logger<T>: Send {
/// Logs an item.
///
/// # Arguments
///
/// * `item` - The item.
    fn log(&mut self, item: T);
/// The logger backend trait.
pub trait LoggerBackend {
/// The logger type.
    type Logger<T>: Logger<T>;
/// Create a new logger.
///
/// # Arguments
///
/// * `epoch` - The epoch.
///
/// # Returns
///
/// The logger.
    fn create<T>(&self, epoch: usize) -> Self::Logger<T>;
/// File logger.
pub struct FileLogger {
impl FileLogger {
/// Create a new file logger.
///
/// # Arguments
///
/// * `path` - The path.
///
/// # Returns
///
/// The file logger.
    pub fn new(path: impl AsRef<Path>) -> Self {
impl<T> Logger<T> for FileLogger
where
    T: std::fmt::Display,
{
    fn log(&mut self, item: T) {
/// In memory logger.
pub struct InMemoryLogger {
impl<T> Logger<T> for InMemoryLogger
where
    T: std::fmt::Display,
{
    fn log(&mut self, item: T) {
const EPOCH_PREFIX: &str = "epoch-";
/// Metric logger.
pub trait MetricLogger: Send {
/// Logs an item.
///
/// # Arguments
///
/// * `item` - The item.
    fn log(&mut self, item: &MetricEntry);
/// Logs an epoch.
///
/// # Arguments
///
/// * `epoch` - The epoch.
    fn end_epoch(&mut self, epoch: usize);
/// Read the logs for an epoch.
    fn read_numeric(&mut self, name: &str, epoch: usize) -> Result<Vec<NumericEntry>, String>;
/// The file metric logger.
pub struct FileMetricLogger {
impl FileMetricLogger {
/// Create a new file metric logger.
///
/// # Arguments
///
/// * `directory` - The directory.
///
/// # Returns
///
/// The file metric logger.
    pub fn new(directory: impl AsRef<Path>) -> Self {
/// Number of epochs recorded.
    pub(crate) fn epochs(&self) -> usize {
    fn epoch_directory(&self, epoch: usize) -> PathBuf {
    fn file_path(&self, name: &str, epoch: usize) -> PathBuf {
    fn create_directory(&self, epoch: usize) {
impl MetricLogger for FileMetricLogger {
    fn log(&mut self, item: &MetricEntry) {
    fn end_epoch(&mut self, epoch: usize) {
    fn read_numeric(&mut self, name: &str, epoch: usize) -> Result<Vec<NumericEntry>, String> {
/// In memory metric logger, useful when testing and debugging.
pub struct InMemoryMetricLogger {
impl InMemoryMetricLogger {
/// Create a new in-memory metric logger.
    pub fn new() -> Self {
impl MetricLogger for InMemoryMetricLogger {
    fn log(&mut self, item: &MetricEntry) {
    fn end_epoch(&mut self, _epoch: usize) {
    fn read_numeric(&mut self, name: &str, epoch: usize) -> Result<Vec<NumericEntry>, String> {
/// The accuracy metric.
pub struct AccuracyMetric<B: Backend> {
/// The [accuracy metric](AccuracyMetric) input type.
pub struct AccuracyInput<B: Backend> {
impl<B: Backend> AccuracyMetric<B> {
/// Creates the metric.
    pub fn new() -> Self {
/// Sets the pad token.
    pub fn with_pad_token(mut self, index: usize) -> Self {
impl<B: Backend> Metric for AccuracyMetric<B> {
    type Input = AccuracyInput<B>;
    fn update(&mut self, input: &AccuracyInput<B>, _metadata: &MetricMetadata) -> MetricEntry {
    fn clear(&mut self) {
    fn name(&self) -> String {
impl<B: Backend> Numeric for AccuracyMetric<B> {
    fn value(&self) -> f64 {
    fn test_accuracy_without_padding() {
// 2
// 1
// 0
// 1
    fn test_accuracy_with_padding() {
// 2
// 1
// 0
// 1
// Predicted padding should not count
// Error on padding should not count
// Error on padding should not count
/// The Area Under the Receiver Operating Characteristic Curve (AUROC, also referred to as [ROC AUC](https://en.wikipedia.org/wiki/Receiver_operating_characteristic)) for binary classification.
pub struct AurocMetric<B: Backend> {
/// The [AUROC metric](AurocMetric) input type.
pub struct AurocInput<B: Backend> {
impl<B: Backend> AurocMetric<B> {
/// Creates the metric.
    pub fn new() -> Self {
    fn binary_auroc(&self, probabilities: &Tensor<B, 1>, targets: &Tensor<B, 1, Int>) -> f64 {
// Early return if we don't have both positive and negative samples
// Calculate AUC components
impl<B: Backend> Metric for AurocMetric<B> {
    type Input = AurocInput<B>;
    fn update(&mut self, input: &AurocInput<B>, _metadata: &MetricMetadata) -> MetricEntry {
    fn clear(&mut self) {
    fn name(&self) -> String {
impl<B: Backend> Numeric for AurocMetric<B> {
    fn value(&self) -> f64 {
    fn test_auroc() {
// High confidence positive
// Low confidence negative
// Low confidence negative
// High confidence positive
// True labels
    fn test_auroc_perfect_separation() {
// Perfect AUC
    fn test_auroc_random() {
// Random predictions
    fn test_auroc_all_one_class() {
// All positives predictions
// All positive class
    fn test_auroc_multiclass_error() {
// More than 2 classes not supported
/// Metric metadata that can be used when computing metrics.
pub struct MetricMetadata {
/// The current progress.
/// The current epoch.
/// The total number of epochs.
/// The current iteration.
/// The current learning rate.
impl MetricMetadata {
/// Fake metric metadata
    pub fn fake() -> Self {
/// Metric trait.
///
/// # Notes
///
/// Implementations should define their own input type only used by the metric.
/// This is important since some conflict may happen when the model output is adapted for each
/// metric's input type.
pub trait Metric: Send + Sync {
/// The input type of the metric.
    type Input;
/// The parameterized name of the metric.
///
/// This should be unique, so avoid using short generic names, prefer using the long name.
///
/// For a metric that can exist at different parameters (e.g., top-k accuracy for different
/// values of k), the name should be unique for each instance.
    fn name(&self) -> String;
/// Update the metric state and returns the current metric entry.
    fn update(&mut self, item: &Self::Input, metadata: &MetricMetadata) -> MetricEntry;
/// Clear the metric state.
    fn clear(&mut self);
/// Adaptor are used to transform types so that they can be used by metrics.
///
/// This should be implemented by a model's output type for all [metric inputs](Metric::Input) that are
/// registered with the [learner builder](crate::learner::LearnerBuilder) .
pub trait Adaptor<T> {
/// Adapt the type to be passed to a [metric](Metric).
    fn adapt(&self) -> T;
impl<T> Adaptor<()> for T {
    fn adapt(&self) {
/// Declare a metric to be numeric.
///
/// This is useful to plot the values of a metric during training.
pub trait Numeric {
/// Returns the numeric value of the metric.
    fn value(&self) -> f64;
/// Data type that contains the current state of a metric at a given time.
pub struct MetricEntry {
/// The name of the metric.
/// The string to be displayed.
/// The string to be saved.
/// Numeric metric entry.
pub enum NumericEntry {
/// Single numeric value.
/// Aggregated numeric (value, number of elements).
impl NumericEntry {
    pub(crate) fn serialize(&self) -> String {
    pub(crate) fn deserialize(entry: &str) -> Result<Self, String> {
// Check for comma separated values
// Numeric value
// Aggregated numeric (value, number of elements)
/// Format a float with the given precision. Will use scientific notation if necessary.
pub fn format_float(float: f64, precision: usize) -> String {
/// Necessary data for classification metrics.
pub struct ClassificationMetricConfig {
/// The prediction decision rule for classification metrics.
pub enum DecisionRule {
/// Consider a class predicted if its probability exceeds the threshold.
/// Consider a class predicted correctly if it is within the top k predicted classes based on scores.
impl Default for DecisionRule {
    fn default() -> Self {
/// The reduction strategy for classification metrics.
pub enum ClassReduction {
/// Computes the statistics over all classes before averaging
/// Computes the statistics independently for each class before averaging
/// Input for confusion statistics error types.
pub struct ConfusionStatsInput<B: Backend> {
/// Sample x Class Non thresholded normalized predictions.
/// Sample x Class one-hot encoded target.
impl<B: Backend> From<ConfusionStatsInput<B>> for (Tensor<B, 2>, Tensor<B, 2, Bool>) {
    fn from(input: ConfusionStatsInput<B>) -> Self {
impl<B: Backend> From<(Tensor<B, 2>, Tensor<B, 2, Bool>)> for ConfusionStatsInput<B> {
    fn from(value: (Tensor<B, 2>, Tensor<B, 2, Bool>)) -> Self {
pub struct ConfusionStats<B: Backend> {
impl<B: Backend> Debug for ConfusionStats<B> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
impl<B: Backend> ConfusionStats<B> {
/// Expects `predictions` to be normalized.
    pub fn new(input: &ConfusionStatsInput<B>, config: &ClassificationMetricConfig) -> Self {
/// sum over samples
    fn aggregate(
        sample_class_mask: Tensor<B, 2, Bool>,
        class_reduction: ClassReduction,
    ) -> Tensor<B, 1> {
    pub fn true_positive(self) -> Tensor<B, 1> {
    pub fn true_negative(self) -> Tensor<B, 1> {
    pub fn false_positive(self) -> Tensor<B, 1> {
    pub fn false_negative(self) -> Tensor<B, 1> {
    pub fn positive(self) -> Tensor<B, 1> {
    pub fn negative(self) -> Tensor<B, 1> {
    pub fn predicted_positive(self) -> Tensor<B, 1> {
    pub fn support(self) -> Tensor<B, 1> {
    pub fn ratio_of_support(self, metric: Tensor<B, 1>) -> Tensor<B, 1> {
    fn top_k_config(
        top_k: NonZeroUsize,
        class_reduction: ClassReduction,
    ) -> ClassificationMetricConfig {
    fn top_k_config_k1_micro() -> ClassificationMetricConfig {
    fn top_k_config_k1_macro() -> ClassificationMetricConfig {
    fn top_k_config_k2_micro() -> ClassificationMetricConfig {
    fn top_k_config_k2_macro() -> ClassificationMetricConfig {
    fn threshold_config(
        threshold: f64,
        class_reduction: ClassReduction,
    ) -> ClassificationMetricConfig {
    fn threshold_config_micro() -> ClassificationMetricConfig {
    fn threshold_config_macro() -> ClassificationMetricConfig {
    fn test_true_positive(
        #[case] classification_type: ClassificationType,
        #[case] config: ClassificationMetricConfig,
        #[case] expected: Vec<i64>,
    ) {
    fn test_true_negative(
        #[case] classification_type: ClassificationType,
        #[case] config: ClassificationMetricConfig,
        #[case] expected: Vec<i64>,
    ) {
    fn test_false_positive(
        #[case] classification_type: ClassificationType,
        #[case] config: ClassificationMetricConfig,
        #[case] expected: Vec<i64>,
    ) {
    fn test_false_negatives(
        #[case] classification_type: ClassificationType,
        #[case] config: ClassificationMetricConfig,
        #[case] expected: Vec<i64>,
    ) {
    fn test_positive(
        #[case] classification_type: ClassificationType,
        #[case] config: ClassificationMetricConfig,
        #[case] expected: Vec<i64>,
    ) {
    fn test_negative(
        #[case] classification_type: ClassificationType,
        #[case] config: ClassificationMetricConfig,
        #[case] expected: Vec<i64>,
    ) {
    fn test_predicted_positive(
        #[case] classification_type: ClassificationType,
        #[case] config: ClassificationMetricConfig,
        #[case] expected: Vec<i64>,
    ) {
/// CPU Temperature metric
/// CPU Temperature in celsius degrees
pub struct CpuTemperature {
impl CpuTemperature {
/// Creates a new CPU temp metric
    pub fn new() -> Self {
impl Default for CpuTemperature {
    fn default() -> Self {
impl Metric for CpuTemperature {
    type Input = ();
    fn update(&mut self, _item: &Self::Input, _metadata: &MetricMetadata) -> MetricEntry {
    fn clear(&mut self) {
    fn name(&self) -> String {
impl Numeric for CpuTemperature {
    fn value(&self) -> f64 {
/// General CPU Usage metric
pub struct CpuUse {
impl CpuUse {
/// Creates a new CPU metric
    pub fn new() -> Self {
    fn refresh(sys: &mut System) -> f64 {
impl Default for CpuUse {
    fn default() -> Self {
impl Metric for CpuUse {
    type Input = ();
    fn update(&mut self, _item: &Self::Input, _metadata: &MetricMetadata) -> MetricEntry {
    fn clear(&mut self) {
    fn name(&self) -> String {
impl Numeric for CpuUse {
    fn value(&self) -> f64 {
/// Track basic cuda infos.
pub struct CudaMetric {
impl CudaMetric {
/// Creates a new metric for CUDA.
    pub fn new() -> Self {
impl Default for CudaMetric {
    fn default() -> Self {
impl Metric for CudaMetric {
    type Input = ();
    fn update(&mut self, _item: &(), _metadata: &MetricMetadata) -> MetricEntry {
    fn clear(&mut self) {
    fn name(&self) -> String {
/// The [F-beta score](https://en.wikipedia.org/wiki/F-score) metric.
///
/// The `beta` parameter represents the ratio of recall importance to precision importance.
/// `beta > 1` gives more weight to recall, while `beta < 1` favors precision.
pub struct FBetaScoreMetric<B: Backend> {
impl<B: Backend> FBetaScoreMetric<B> {
/// F-beta score metric for binary classification.
///
/// # Arguments
///
/// * `beta` - Positive real factor to weight recall's importance.
/// * `threshold` - The threshold to transform a probability into a binary prediction.
    pub fn binary(beta: f64, threshold: f64) -> Self {
// binary classification results are the same independently of class_reduction
/// F-beta score metric for multiclass classification.
///
/// # Arguments
///
/// * `beta` - Positive real factor to weight recall's importance.
/// * `top_k` - The number of highest predictions considered to find the correct label (typically `1`).
/// * `class_reduction` - [Class reduction](ClassReduction) type.
    pub fn multiclass(beta: f64, top_k: usize, class_reduction: ClassReduction) -> Self {
/// F-beta score metric for multi-label classification.
///
/// # Arguments
///
/// * `beta` - Positive real factor to weight recall's importance.
/// * `threshold` - The threshold to transform a probability into a binary prediction.
/// * `class_reduction` - [Class reduction](ClassReduction) type.
    pub fn multilabel(beta: f64, threshold: f64, class_reduction: ClassReduction) -> Self {
    fn class_average(&self, mut aggregated_metric: Tensor<B, 1>) -> f64 {
impl<B: Backend> Metric for FBetaScoreMetric<B> {
    type Input = ConfusionStatsInput<B>;
    fn update(&mut self, input: &Self::Input, _metadata: &MetricMetadata) -> MetricEntry {
    fn clear(&mut self) {
    fn name(&self) -> String {
// "FBetaScore (0.5) @ TopK(1) [Macro]"
impl<B: Backend> Numeric for FBetaScoreMetric<B> {
    fn value(&self) -> f64 {
    fn test_binary_fscore(#[case] beta: f64, #[case] threshold: f64, #[case] expected: f64) {
    fn test_multiclass_fscore(
        #[case] beta: f64,
        #[case] class_reduction: ClassReduction,
        #[case] top_k: usize,
        #[case] expected: f64,
    ) {
    fn test_multilabel_fscore(
        #[case] beta: f64,
        #[case] class_reduction: ClassReduction,
        #[case] threshold: f64,
        #[case] expected: f64,
    ) {
    fn test_parameterized_unique_name() {
/// The hamming score, sometimes referred to as multi-label or label-based accuracy.
pub struct HammingScore<B: Backend> {
/// The [hamming score](HammingScore) input type.
pub struct HammingScoreInput<B: Backend> {
impl<B: Backend> HammingScore<B> {
/// Creates the metric.
    pub fn new() -> Self {
/// Sets the threshold.
    pub fn with_threshold(mut self, threshold: f32) -> Self {
/// Sets the sigmoid activation function usage.
    pub fn with_sigmoid(mut self, sigmoid: bool) -> Self {
impl<B: Backend> Default for HammingScore<B> {
/// Creates a new metric instance with default values.
    fn default() -> Self {
impl<B: Backend> Metric for HammingScore<B> {
    type Input = HammingScoreInput<B>;
    fn update(&mut self, input: &HammingScoreInput<B>, _metadata: &MetricMetadata) -> MetricEntry {
    fn clear(&mut self) {
    fn name(&self) -> String {
impl<B: Backend> Numeric for HammingScore<B> {
    fn value(&self) -> f64 {
    fn test_hamming_score() {
// with x > 0.5: [0, 1, 0, 1, 1]
//               [0, 0, 0, 1, 1]
//               [0, 0, 1, 0, 1]
//               [0, 0, 1, 0, 0]
// Invert all targets: y = (1 - y)
// invert targets (1 - y)
// Invert 5 target values -> 1 - (5/20) = 0.75
// invert targets (1 - y)
    fn test_parameterized_unique_name() {
/// The loss metric.
pub struct IterationSpeedMetric {
impl IterationSpeedMetric {
/// Create the metric.
    pub fn new() -> Self {
impl Metric for IterationSpeedMetric {
    type Input = ();
    fn update(&mut self, _: &Self::Input, metadata: &MetricMetadata) -> MetricEntry {
    fn clear(&mut self) {
    fn name(&self) -> String {
impl Numeric for IterationSpeedMetric {
    fn value(&self) -> f64 {
/// Track the learning rate across iterations.
pub struct LearningRateMetric {
impl LearningRateMetric {
/// Creates a new learning rate metric.
    pub fn new() -> Self {
impl Default for LearningRateMetric {
    fn default() -> Self {
impl Metric for LearningRateMetric {
    type Input = ();
    fn update(&mut self, _item: &(), metadata: &MetricMetadata) -> MetricEntry {
    fn clear(&mut self) {
    fn name(&self) -> String {
impl Numeric for LearningRateMetric {
    fn value(&self) -> f64 {
/// The loss metric.
pub struct LossMetric<B: Backend> {
/// The [loss metric](LossMetric) input type.
pub struct LossInput<B: Backend> {
impl<B: Backend> LossMetric<B> {
/// Create the metric.
    pub fn new() -> Self {
impl<B: Backend> Metric for LossMetric<B> {
    type Input = LossInput<B>;
    fn update(&mut self, loss: &Self::Input, _metadata: &MetricMetadata) -> MetricEntry {
    fn clear(&mut self) {
    fn name(&self) -> String {
impl<B: Backend> Numeric for LossMetric<B> {
    fn value(&self) -> f64 {
/// RAM use metric
/// Memory information
pub struct CpuMemory {
impl CpuMemory {
/// Creates a new memory metric
    pub fn new() -> Self {
    fn refresh(&mut self) {
// bytes of RAM available
// bytes of RAM in use
impl Default for CpuMemory {
    fn default() -> Self {
impl Metric for CpuMemory {
    type Input = ();
    fn update(&mut self, _item: &Self::Input, _metadata: &MetricMetadata) -> MetricEntry {
    fn clear(&mut self) {
    fn name(&self) -> String {
impl Numeric for CpuMemory {
    fn value(&self) -> f64 {
fn bytes2gb(bytes: u64) -> f64 {
/// State module.
/// Module responsible to save and exposes data collected during training.
// System metrics
// Training metrics
// Expose `ItemLazy` so it can be implemented for custom types
/// The Precision Metric
pub struct PrecisionMetric<B: Backend> {
impl<B: Backend> PrecisionMetric<B> {
/// Precision metric for binary classification.
///
/// # Arguments
///
/// * `threshold` - The threshold to transform a probability into a binary prediction.
    pub fn binary(threshold: f64) -> Self {
// binary classification results are the same independently of class_reduction
/// Precision metric for multiclass classification.
///
/// # Arguments
///
/// * `top_k` - The number of highest predictions considered to find the correct label (typically `1`).
/// * `class_reduction` - [Class reduction](ClassReduction) type.
    pub fn multiclass(top_k: usize, class_reduction: ClassReduction) -> Self {
/// Precision metric for multi-label classification.
///
/// # Arguments
///
/// * `threshold` - The threshold to transform a probability into a binary value.
/// * `class_reduction` - [Class reduction](ClassReduction) type.
    pub fn multilabel(threshold: f64, class_reduction: ClassReduction) -> Self {
    fn class_average(&self, mut aggregated_metric: Tensor<B, 1>) -> f64 {
impl<B: Backend> Metric for PrecisionMetric<B> {
    type Input = ConfusionStatsInput<B>;
    fn update(&mut self, input: &Self::Input, _metadata: &MetricMetadata) -> MetricEntry {
    fn clear(&mut self) {
    fn name(&self) -> String {
// "Precision @ Threshold(0.5) [Macro]"
impl<B: Backend> Numeric for PrecisionMetric<B> {
    fn value(&self) -> f64 {
    fn test_binary_precision(#[case] threshold: f64, #[case] expected: f64) {
    fn test_multiclass_precision(
        #[case] class_reduction: ClassReduction,
        #[case] top_k: usize,
        #[case] expected: f64,
    ) {
    fn test_multilabel_precision(
        #[case] class_reduction: ClassReduction,
        #[case] threshold: f64,
        #[case] expected: f64,
    ) {
    fn test_parameterized_unique_name() {
pub struct AsyncProcessor<P: EventProcessor> {
struct Worker<P: EventProcessor> {
impl<P: EventProcessor + 'static> Worker<P> {
    pub fn start(processor: P, rec: Receiver<Message<P>>) {
impl<P: EventProcessor + 'static> AsyncProcessor<P> {
    pub fn new(processor: P) -> Self {
enum Message<P: EventProcessor> {
impl<P: EventProcessor> EventProcessor for AsyncProcessor<P> {
    type ItemTrain = P::ItemTrain;
    type ItemValid = P::ItemValid;
    fn process_train(&mut self, event: Event<Self::ItemTrain>) {
    fn process_valid(&mut self, event: Event<Self::ItemValid>) {
/// Event happening during the training/validation process.
pub enum Event<T> {
/// Signal that an item have been processed.
/// Signal the end of an epoch.
/// Signal the end of the process (e.g., training end).
/// Items that are lazy are not ready to be processed by metrics.
///
/// We want to sync them on a different thread to avoid blocking training.
pub trait ItemLazy: Send {
/// Item that is properly synced and ready to be processed by metrics.
    type ItemSync: Send;
/// Sync the item.
    fn sync(self) -> Self::ItemSync;
/// Process events happening during training and validation.
pub trait EventProcessor: Send {
/// The training item.
    type ItemTrain: ItemLazy;
/// The validation item.
    type ItemValid: ItemLazy;
/// Collect a training event.
    fn process_train(&mut self, event: Event<Self::ItemTrain>);
/// Collect a validation event.
    fn process_valid(&mut self, event: Event<Self::ItemValid>);
/// A learner item.
pub struct LearnerItem<T> {
/// The item.
/// The progress.
/// The epoch.
/// The total number of epochs.
/// The iteration.
/// The learning rate.
impl<T: ItemLazy> ItemLazy for LearnerItem<T> {
    type ItemSync = LearnerItem<T::ItemSync>;
    fn sync(self) -> Self::ItemSync {
/// An [event processor](EventProcessor) that handles:
///   - Computing and storing metrics in an [event store](crate::metric::store::EventStore).
///   - Render metrics using a [metrics renderer](MetricsRenderer).
pub struct FullEventProcessor<T: ItemLazy, V: ItemLazy> {
impl<T: ItemLazy, V: ItemLazy> FullEventProcessor<T, V> {
    pub(crate) fn new(
        metrics: Metrics<T, V>,
        renderer: Box<dyn MetricsRenderer>,
        store: Arc<EventStoreClient>,
    ) -> Self {
impl<T: ItemLazy, V: ItemLazy> EventProcessor for FullEventProcessor<T, V> {
    type ItemTrain = T;
    type ItemValid = V;
    fn process_train(&mut self, event: Event<Self::ItemTrain>) {
    fn process_valid(&mut self, event: Event<Self::ItemValid>) {
// no-op for now
pub(crate) struct Metrics<T: ItemLazy, V: ItemLazy> {
impl<T: ItemLazy, V: ItemLazy> Default for Metrics<T, V> {
    fn default() -> Self {
impl<T: ItemLazy, V: ItemLazy> Metrics<T, V> {
/// Register a training metric.
    pub(crate) fn register_train_metric<Me: Metric + 'static>(&mut self, metric: Me)
    where
        T::ItemSync: Adaptor<Me::Input> + 'static,
    {
/// Register a validation metric.
    pub(crate) fn register_valid_metric<Me: Metric + 'static>(&mut self, metric: Me)
    where
        V::ItemSync: Adaptor<Me::Input> + 'static,
    {
/// Register a numeric training metric.
    pub(crate) fn register_train_metric_numeric<Me: Metric + Numeric + 'static>(
        &mut self,
        metric: Me,
    ) where
        T::ItemSync: Adaptor<Me::Input> + 'static,
    {
/// Register a numeric validation metric.
    pub(crate) fn register_valid_metric_numeric<Me: Metric + Numeric + 'static>(
        &mut self,
        metric: Me,
    ) where
        V::ItemSync: Adaptor<Me::Input> + 'static,
    {
/// Update the training information from the training item.
    pub(crate) fn update_train(
        &mut self,
        item: &LearnerItem<T::ItemSync>,
        metadata: &MetricMetadata,
    ) -> MetricsUpdate {
/// Update the training information from the validation item.
    pub(crate) fn update_valid(
        &mut self,
        item: &LearnerItem<V::ItemSync>,
        metadata: &MetricMetadata,
    ) -> MetricsUpdate {
/// Signal the end of a training epoch.
    pub(crate) fn end_epoch_train(&mut self) {
/// Signal the end of a validation epoch.
    pub(crate) fn end_epoch_valid(&mut self) {
impl<T> From<&LearnerItem<T>> for TrainingProgress {
    fn from(item: &LearnerItem<T>) -> Self {
impl<T> From<&LearnerItem<T>> for MetricMetadata {
    fn from(item: &LearnerItem<T>) -> Self {
trait NumericMetricUpdater<T>: Send + Sync {
    fn update(&mut self, item: &LearnerItem<T>, metadata: &MetricMetadata) -> (MetricEntry, f64);
    fn clear(&mut self);
trait MetricUpdater<T>: Send + Sync {
    fn update(&mut self, item: &LearnerItem<T>, metadata: &MetricMetadata) -> MetricEntry;
    fn clear(&mut self);
struct MetricWrapper<M> {
impl<T, M> NumericMetricUpdater<T> for MetricWrapper<M>
where
    T: 'static,
    M: Metric + Numeric + 'static,
    T: Adaptor<M::Input>,
{
    fn update(&mut self, item: &LearnerItem<T>, metadata: &MetricMetadata) -> (MetricEntry, f64) {
    fn clear(&mut self) {
impl<T, M> MetricUpdater<T> for MetricWrapper<M>
where
    T: 'static,
    M: Metric + 'static,
    T: Adaptor<M::Input>,
{
    fn update(&mut self, item: &LearnerItem<T>, metadata: &MetricMetadata) -> MetricEntry {
    fn clear(&mut self) {
/// An [event processor](EventProcessor) that handles:
///   - Computing and storing metrics in an [event store](crate::metric::store::EventStore).
pub(crate) struct MinimalEventProcessor<T: ItemLazy, V: ItemLazy> {
impl<T: ItemLazy, V: ItemLazy> EventProcessor for MinimalEventProcessor<T, V> {
    type ItemTrain = T;
    type ItemValid = V;
    fn process_train(&mut self, event: Event<Self::ItemTrain>) {
// no-op for now
    fn process_valid(&mut self, event: Event<Self::ItemValid>) {
// no-op for now
    impl ItemLazy for f64 {
        type ItemSync = f64;
        fn sync(self) -> Self::ItemSync {
    impl<B: Backend> Adaptor<LossInput<B>> for f64 {
        fn adapt(&self) -> LossInput<B> {
    pub(crate) fn process_train(
        processor: &mut MinimalEventProcessor<f64, f64>,
        value: f64,
        epoch: usize,
    ) {
    pub(crate) fn end_epoch(processor: &mut MinimalEventProcessor<f64, f64>, epoch: usize) {
///The Recall Metric
pub struct RecallMetric<B: Backend> {
impl<B: Backend> RecallMetric<B> {
/// Recall metric for binary classification.
///
/// # Arguments
///
/// * `threshold` - The threshold to transform a probability into a binary prediction.
    pub fn binary(threshold: f64) -> Self {
// binary classification results are the same independently of class_reduction
/// Recall metric for multiclass classification.
///
/// # Arguments
///
/// * `top_k` - The number of highest predictions considered to find the correct label (typically `1`).
/// * `class_reduction` - [Class reduction](ClassReduction) type.
    pub fn multiclass(top_k: usize, class_reduction: ClassReduction) -> Self {
/// Recall metric for multi-label classification.
///
/// # Arguments
///
/// * `threshold` - The threshold to transform a probability into a binary prediction.
/// * `class_reduction` - [Class reduction](ClassReduction) type.
    pub fn multilabel(threshold: f64, class_reduction: ClassReduction) -> Self {
    fn class_average(&self, mut aggregated_metric: Tensor<B, 1>) -> f64 {
impl<B: Backend> Metric for RecallMetric<B> {
    type Input = ConfusionStatsInput<B>;
    fn update(&mut self, input: &Self::Input, _metadata: &MetricMetadata) -> MetricEntry {
    fn clear(&mut self) {
    fn name(&self) -> String {
// "Recall @ Threshold(0.5) [Macro]"
impl<B: Backend> Numeric for RecallMetric<B> {
    fn value(&self) -> f64 {
    fn test_binary_recall(#[case] threshold: f64, #[case] expected: f64) {
    fn test_multiclass_recall(
        #[case] class_reduction: ClassReduction,
        #[case] top_k: usize,
        #[case] expected: f64,
    ) {
    fn test_multilabel_recall(
        #[case] class_reduction: ClassReduction,
        #[case] threshold: f64,
        #[case] expected: f64,
    ) {
    fn test_parameterized_unique_name() {
/// Useful utility to implement numeric metrics.
///
/// # Notes
///
/// The numeric metric store values inside floats.
/// Even if some metric are integers, their mean are floats.
pub struct NumericMetricState {
/// Formatting options for the [numeric metric state](NumericMetricState).
pub struct FormatOptions {
impl FormatOptions {
/// Create the [formatting options](FormatOptions) with a name.
    pub fn new<S: Into<String>>(name: S) -> Self {
/// Specify the metric unit.
    pub fn unit(mut self, unit: &str) -> Self {
/// Specify the floating point precision.
    pub fn precision(mut self, precision: usize) -> Self {
impl NumericMetricState {
/// Create a new [numeric metric state](NumericMetricState).
    pub fn new() -> Self {
/// Reset the state.
    pub fn reset(&mut self) {
/// Update the state.
    pub fn update(&mut self, value: f64, batch_size: usize, format: FormatOptions) -> MetricEntry {
// Numeric metric state is an aggregated value
impl Numeric for NumericMetricState {
    fn value(&self) -> f64 {
impl Default for NumericMetricState {
    fn default() -> Self {
/// Type that can be used to fetch and use numeric metric aggregates.
pub(crate) struct NumericMetricsAggregate {
struct Key {
impl NumericMetricsAggregate {
    pub(crate) fn aggregate(
        &mut self,
        name: &str,
        epoch: usize,
        aggregate: Aggregate,
        loggers: &mut [Box<dyn MetricLogger>],
    ) -> Option<f64> {
// Accurately compute the aggregated value based on the *actual* number of points
// since not all mini-batches are guaranteed to have the specified batch size
// Right now the mean is the only aggregate available, so we can assume that the sum
// of an entry corresponds to (value * number of elements)
    pub(crate) fn find_epoch(
        &mut self,
        name: &str,
        aggregate: Aggregate,
        direction: Direction,
        loggers: &mut [Box<dyn MetricLogger>],
    ) -> Option<usize> {
    struct TestLogger {
    const NAME: &str = "test-logger";
    impl TestLogger {
        fn new() -> Self {
        fn log(&mut self, num: f64) {
        fn new_epoch(&mut self) {
    fn should_find_epoch() {
// Epoch 1
// Epoch 1
// Epoch 2
// Epoch 2
// Epoch 3
    fn should_aggregate_numeric_entry() {
// Epoch 1
// (1.5 + 1.0) / 2 = 2.5 / 2
// Average should be (0.5 + 1.25 * 2) / 3 = 1.0, not (0.5 + 1.25) / 2 = 0.875
/// Event happening during the training/validation process.
pub enum Event {
/// Signal that metrics have been updated.
/// Signal the end of an epoch.
/// Contains all metric information.
pub struct MetricsUpdate {
/// Metrics information related to non-numeric metrics.
/// Metrics information related to numeric metrics.
/// Defines how training and validation events are collected and searched.
///
/// This trait also exposes methods that uses the collected data to compute useful information.
pub trait EventStore: Send {
/// Collect a training/validation event.
    fn add_event(&mut self, event: Event, split: Split);
/// Find the epoch following the given criteria from the collected data.
    fn find_epoch(
        &mut self,
        name: &str,
        aggregate: Aggregate,
        direction: Direction,
        split: Split,
    ) -> Option<usize>;
/// Find the metric value for the current epoch following the given criteria.
    fn find_metric(
        &mut self,
        name: &str,
        epoch: usize,
        aggregate: Aggregate,
        split: Split,
    ) -> Option<f64>;
/// How to aggregate the metric.
pub enum Aggregate {
/// Compute the average.
/// The split to use.
pub enum Split {
/// The training split.
/// The validation split.
/// The direction of the query.
pub enum Direction {
/// Lower is better.
/// Higher is better.
/// Type that allows to communicate with an [event store](EventStore).
pub struct EventStoreClient {
impl EventStoreClient {
/// Create a new [event store](EventStore) client.
    pub(crate) fn new<C>(store: C) -> Self
    where
        C: EventStore + 'static,
    {
impl EventStoreClient {
/// Add a training event to the [event store](EventStore).
    pub(crate) fn add_event_train(&self, event: Event) {
/// Add a validation event to the [event store](EventStore).
    pub(crate) fn add_event_valid(&self, event: Event) {
/// Find the epoch following the given criteria from the collected data.
    pub fn find_epoch(
        &self,
        name: &str,
        aggregate: Aggregate,
        direction: Direction,
        split: Split,
    ) -> Option<usize> {
/// Find the metric value for the current epoch following the given criteria.
    pub fn find_metric(
        &self,
        name: &str,
        epoch: usize,
        aggregate: Aggregate,
        split: Split,
    ) -> Option<f64> {
struct WorkerThread<S> {
impl<C> WorkerThread<C>
where
    C: EventStore,
{
    fn run(mut self) {
enum Message {
impl Drop for EventStoreClient {
    fn drop(&mut self) {
pub(crate) struct LogEventStore {
impl EventStore for LogEventStore {
    fn add_event(&mut self, event: Event, split: Split) {
    fn find_epoch(
        &mut self,
        name: &str,
        aggregate: Aggregate,
        direction: Direction,
        split: Split,
    ) -> Option<usize> {
    fn find_metric(
        &mut self,
        name: &str,
        epoch: usize,
        aggregate: Aggregate,
        split: Split,
    ) -> Option<f64> {
impl LogEventStore {
/// Register a logger for training metrics.
    pub(crate) fn register_logger_train<ML: MetricLogger + 'static>(&mut self, logger: ML) {
/// Register a logger for validation metrics.
    pub(crate) fn register_logger_valid<ML: MetricLogger + 'static>(&mut self, logger: ML) {
/// The Top-K accuracy metric.
///
/// For K=1, this is equivalent to the [accuracy metric](`super::acc::AccuracyMetric`).
pub struct TopKAccuracyMetric<B: Backend> {
/// If specified, targets equal to this value will be considered padding and will not count
/// towards the metric
/// The [top-k accuracy metric](TopKAccuracyMetric) input type.
pub struct TopKAccuracyInput<B: Backend> {
/// The outputs (batch_size, num_classes)
/// The labels (batch_size)
impl<B: Backend> TopKAccuracyMetric<B> {
/// Creates the metric.
    pub fn new(k: usize) -> Self {
/// Sets the pad token.
    pub fn with_pad_token(mut self, index: usize) -> Self {
impl<B: Backend> Metric for TopKAccuracyMetric<B> {
    type Input = TopKAccuracyInput<B>;
    fn update(&mut self, input: &TopKAccuracyInput<B>, _metadata: &MetricMetadata) -> MetricEntry {
// we ignore the samples where the target is equal to the pad token
    fn clear(&mut self) {
    fn name(&self) -> String {
impl<B: Backend> Numeric for TopKAccuracyMetric<B> {
    fn value(&self) -> f64 {
    fn test_accuracy_without_padding() {
// 2, 1
// 1, 0
// 0, 2
// 1, 0
    fn test_accuracy_with_padding() {
// 2, 1
// 1, 0
// 0, 2
// 1, 0
// Predicted padding should not count
// Error on padding should not count
// Error on padding should not count
    fn test_parameterized_unique_name() {
/// Trait for rendering metrics.
pub trait MetricsRenderer: Send + Sync {
/// Updates the training metric state.
///
/// # Arguments
///
/// * `state` - The metric state.
    fn update_train(&mut self, state: MetricState);
/// Updates the validation metric state.
///
/// # Arguments
///
/// * `state` - The metric state.
    fn update_valid(&mut self, state: MetricState);
/// Renders the training progress.
///
/// # Arguments
///
/// * `item` - The training progress.
    fn render_train(&mut self, item: TrainingProgress);
/// Renders the validation progress.
///
/// # Arguments
///
/// * `item` - The validation progress.
    fn render_valid(&mut self, item: TrainingProgress);
/// Callback method invoked when training ends, whether it
/// completed successfully or was interrupted.
///
/// # Returns
///
/// A result indicating whether the end-of-training actions were successful.
    fn on_train_end(&mut self) -> Result<(), Box<dyn core::error::Error>> {
/// The state of a metric.
pub enum MetricState {
/// A generic metric.
/// A numeric metric.
/// Training progress.
pub struct TrainingProgress {
/// The progress.
/// The epoch.
/// The total number of epochs.
/// The iteration.
impl TrainingProgress {
/// Creates a new empty training progress.
    pub fn none() -> Self {
/// A simple renderer for when the cli feature is not enabled.
pub struct CliMetricsRenderer;
impl CliMetricsRenderer {
/// Create a new instance.
    pub fn new() -> Self {
impl MetricsRenderer for CliMetricsRenderer {
    fn update_train(&mut self, _state: MetricState) {
    fn update_valid(&mut self, _state: MetricState) {
    fn render_train(&mut self, item: TrainingProgress) {
    fn render_valid(&mut self, item: TrainingProgress) {
/// The tui renderer
/// Return the default metrics renderer.
///
/// This can be either:
///   - `TuiMetricsRenderer`, when the `tui` feature is enabled and `stdout` is
///     a terminal, or
///   - `CliMetricsRenderer`, when the `tui` feature is not enabled, or `stdout`
///     is not a terminal.
pub(crate) fn default_renderer(
    interuptor: TrainingInterrupter,
    checkpoint: Option<usize>,
) -> Box<dyn MetricsRenderer> {
pub(crate) struct MetricsView<'a> {
impl MetricsView<'_> {
    pub(crate) fn render(self, frame: &mut TerminalFrame<'_>, size: Rect) {
/// Controls view.
pub(crate) struct ControlsView;
impl ControlsView {
/// Render the view.
    pub(crate) fn render(self, frame: &mut TerminalFrame<'_>, size: Rect) {
/// A plot that shows the full history at a reduced resolution.
pub(crate) struct FullHistoryPlot {
struct FullHistoryPoints {
impl FullHistoryPlot {
/// Create a new history plot.
    pub(crate) fn new(max_samples: usize) -> Self {
/// Update the maximum amount of sample to display for the validation points.
///
/// This is necessary if we want the validation line to have the same point density as the
/// training line.
    pub(crate) fn update_max_sample_valid(&mut self, ratio_train: f64) {
/// Register a training data point.
    pub(crate) fn push_train(&mut self, data: f64) {
/// Register a validation data point.
    pub(crate) fn push_valid(&mut self, data: f64) {
/// Create the training and validation datasets from the data points.
    pub(crate) fn datasets(&self) -> Vec<Dataset<'_>> {
    fn next_x(&mut self) -> f64 {
    fn update_bounds(&mut self) {
impl FullHistoryPoints {
    fn new(max_samples: usize) -> Self {
    fn push(&mut self, (x, y): (f64, f64)) {
/// We keep only half the points and we double the step size.
///
/// This ensure that we have the same amount of points across the X axis.
    fn resize(&mut self) {
    fn dataset<'a>(&'a self, name: &'a str, color: Color) -> Dataset<'a> {
    fn is_empty(&self) -> bool {
    fn test_points() {
/// 1000 seems to be required to see some improvement.
const MAX_NUM_SAMPLES_RECENT: usize = 1000;
/// 250 seems to be the right resolution when plotting all history.
/// Otherwise, there is too much points and the lines arent't smooth enough.
const MAX_NUM_SAMPLES_FULL: usize = 250;
/// Numeric metrics state that handles creating plots.
pub(crate) struct NumericMetricsState {
/// The kind of plot to display.
pub(crate) enum PlotKind {
/// Display the full history of the metric with reduced resolution.
/// Display only the recent history of the metric, but with more resolution.
impl NumericMetricsState {
/// Register a new training value for the metric with the given name.
    pub(crate) fn push_train(&mut self, name: String, data: f64) {
/// Register a new validation value for the metric with the given name.
    pub(crate) fn push_valid(&mut self, key: String, data: f64) {
/// Update the state with the training progress.
    pub(crate) fn update_progress_train(&mut self, progress: &TrainingProgress) {
/// Update the state with the validation progress.
    pub(crate) fn update_progress_valid(&mut self, progress: &TrainingProgress) {
/// Create a view to display the numeric metrics.
    pub(crate) fn view(&self) -> NumericMetricView<'_> {
/// Handle the current event.
    pub(crate) fn on_event(&mut self, event: &Event) {
// Fix the double toggle on Windows.
    fn switch_kind(&mut self) {
    fn next_metric(&mut self) {
    fn previous_metric(&mut self) {
    fn chart<'a>(&'a self) -> Chart<'a> {
pub(crate) enum NumericMetricView<'a> {
impl NumericMetricView<'_> {
    pub(crate) fn render(self, frame: &mut TerminalFrame<'_>, size: Rect) {
pub(crate) struct TextMetricsState {
pub(crate) struct MetricData {
impl TextMetricsState {
    pub(crate) fn update_train(&mut self, metric: MetricEntry) {
    pub(crate) fn update_valid(&mut self, metric: MetricEntry) {
    pub(crate) fn view(&self) -> TextMetricView {
pub(crate) struct TextMetricView {
    lines: Vec<Vec<Span<'static>>>,
}

impl TextMetricView {
    fn new(names: &[String], data: &HashMap<String, MetricData>) -> Self {
    pub(crate) fn render(self, frame: &mut TerminalFrame<'_>, size: Rect) {
const AXIS_TITLE_PRECISION: usize = 2;
/// The data describing both X and Y axes.
pub(crate) struct PlotAxes {
impl Default for PlotAxes {
    fn default() -> Self {
impl PlotAxes {
/// Update the bounds based on the min max of each X and Y axes with both train and valid data.
    pub(crate) fn update_bounds(
        &mut self,
        (x_train_min, x_train_max): (f64, f64),
        (x_valid_min, x_valid_max): (f64, f64),
        (y_train_min, y_train_max): (f64, f64),
        (y_valid_min, y_valid_max): (f64, f64),
    ) {
// We know x are integers.
/// Popup callback function.
pub(crate) trait CallbackFn: Send + Sync {
/// Call the function and return if the popup state should be reset.
    fn call(&self) -> bool;
/// Popup callback.
pub(crate) struct Callback {
impl Callback {
/// Create a new popup.
    pub(crate) fn new<T, D, C>(title: T, description: D, trigger: char, callback: C) -> Self
    where
        T: Into<String>,
        D: Into<String>,
        C: CallbackFn + 'static,
    {
/// Popup state.
pub(crate) enum PopupState {
impl PopupState {
/// If the popup is empty.
    pub(crate) fn is_empty(&self) -> bool {
/// Handle popup events.
    pub(crate) fn on_event(&mut self, event: &Event) {
/// Create the popup view.
    pub(crate) fn view(&self) -> Option<PopupView<'_>> {
pub(crate) struct PopupView<'a> {
impl<'a> PopupView<'a> {
/// Render the view.
    pub(crate) fn render<'b>(&'a self, frame: &mut TerminalFrame<'b>, size: Rect) {
/// The percent represents the amount of space that will be taken by each side.
fn centered_percent(percent: u16, size: Rect, direction: Direction) -> Rect {
/// Simple progress bar for the training.
///
/// We currently ignore the time taken for the validation part.
pub(crate) struct ProgressBarState {
// Progress for total training.
const MINUTE: u64 = 60;
const HOUR: u64 = 60 * 60;
const DAY: u64 = 24 * 60 * 60;
impl ProgressBarState {
    pub fn new(checkpoint: Option<usize>) -> Self {
/// Update the training progress.
    pub(crate) fn update_train(&mut self, progress: &TrainingProgress) {
/// Update the validation progress.
    pub(crate) fn update_valid(&mut self, _progress: &TrainingProgress) {
// We don't use the validation for the progress yet.
/// Create a view for the current progress.
    pub(crate) fn view(&self) -> ProgressBarView {
        const NO_ETA: &str = "---";
pub(crate) struct ProgressBarView {
impl ProgressBarView {
/// Render the view.
    pub(crate) fn render(self, frame: &mut TerminalFrame<'_>, size: Rect) {
// Empty space
struct ProgressEstimate {
impl ProgressEstimate {
    fn new() -> Self {
    fn secs(&self) -> Option<u64> {
    fn update(&mut self, progress: &TrainingProgress, starting_epoch: usize) {
        const WARMUP_NUM_ITERATION: usize = 10;
// When the training has started since 30 seconds.
// When the training has started since at least 10 seconds and completed 10 iterations.
    fn init(&mut self, progress: &TrainingProgress, starting_epoch: usize) {
fn calculate_progress(
    progress: &TrainingProgress,
    starting_epoch: usize,
    ignore_num_items: usize,
) -> f64 {
fn format_eta(eta_secs: u64) -> String {
    fn test_format_eta() {
    fn calculate_progress_for_eta() {
// Two epochs remaining while the first is half done.
    fn calculate_progress_for_eta_with_warmup() {
// Two epochs remaining while the first is half done.
const FACTOR_BEFORE_RESIZE: usize = 2;
/// A plot that shows the recent history at full resolution.
pub(crate) struct RecentHistoryPlot {
struct RecentHistoryPoints {
impl RecentHistoryPlot {
    pub(crate) fn new(max_samples: usize) -> Self {
    pub(crate) fn push_train(&mut self, data: f64) {
    pub(crate) fn push_valid(&mut self, data: f64) {
    pub(crate) fn datasets(&self) -> Vec<Dataset<'_>> {
    fn x(&mut self) -> (f64, f64) {
    fn update_bounds(&mut self) {
impl RecentHistoryPoints {
    fn new(max_samples: usize) -> Self {
    fn num_visible_points(&self) -> usize {
    fn push(&mut self, (x, y): (f64, f64)) {
    fn update_cursor(&mut self, min_x: f64) {
    fn slice(&self) -> &[(f64, f64)] {
    fn calculate_max_y(&self) -> f64 {
    fn calculate_min_y(&self) -> f64 {
    fn resize(&mut self) {
    fn dataset<'a>(&'a self, name: &'a str, color: Color) -> Dataset<'a> {
    fn test_push_update_bounds_max_y() {
    fn test_push_update_bounds_min_y() {
/// The current terminal backend.
pub(crate) type TerminalBackend = CrosstermBackend<Stdout>;
/// The current terminal frame.
pub(crate) type TerminalFrame<'a> = ratatui::Frame<'a>;
type PanicHook = Box<dyn Fn(&std::panic::PanicHookInfo<'_>) + 'static + Sync + Send>;
const MAX_REFRESH_RATE_MILLIS: u64 = 100;
/// The terminal UI metrics renderer.
pub struct TuiMetricsRenderer {
impl MetricsRenderer for TuiMetricsRenderer {
    fn update_train(&mut self, state: MetricState) {
    fn update_valid(&mut self, state: MetricState) {
    fn render_train(&mut self, item: TrainingProgress) {
    fn render_valid(&mut self, item: TrainingProgress) {
    fn on_train_end(&mut self) -> Result<(), Box<dyn Error>> {
impl TuiMetricsRenderer {
/// Create a new terminal UI renderer.
    pub fn new(interuptor: TrainingInterrupter, checkpoint: Option<usize>) -> Self {
// Reset the terminal to raw mode on panic before running the panic handler
// This prevents that the panic message is not visible for the user.
/// Set the renderer to persistent mode.
    pub fn persistent(mut self) -> Self {
    fn render(&mut self) -> Result<(), Box<dyn Error>> {
    fn draw(&mut self) -> Result<(), Box<dyn Error>> {
    fn handle_events(&mut self) -> Result<(), Box<dyn Error>> {
/// Show the training status with various information.
pub(crate) struct StatusState {
enum Mode {
impl Default for StatusState {
    fn default() -> Self {
impl StatusState {
/// Update the training information.
    pub(crate) fn update_train(&mut self, progress: TrainingProgress) {
/// Update the validation information.
    pub(crate) fn update_valid(&mut self, progress: TrainingProgress) {
/// Create a view.
    pub(crate) fn view(&self) -> StatusView {
pub(crate) struct StatusView {
    lines: Vec<Vec<Span<'static>>>,
}

impl StatusView {
    fn new(progress: &TrainingProgress, mode: &Mode) -> Self {
    pub(crate) fn render(self, frame: &mut TerminalFrame<'_>, size: Rect) {
// Vulkan and WebGpu would have conflicting type names
/// Tensor backend that uses the wgpu crate for executing GPU compute shaders.
///
/// This backend can target multiple graphics APIs, including:
///   - [Vulkan][crate::graphics::Vulkan] on Linux, Windows, and Android.
///   - [OpenGL](crate::graphics::OpenGl) on Linux, Windows, and Android.
///   - [DirectX 12](crate::graphics::Dx12) on Windows.
///   - [Metal][crate::graphics::Metal] on Apple hardware.
///   - [WebGPU](crate::graphics::WebGpu) on supported browsers and `wasm` runtimes.
///
/// To configure the wgpu backend, eg. to select what graphics API to use or what memory strategy to use,
/// you have to manually initialize the runtime. For example:
///
/// ```rust, ignore
/// fn custom_init() {
///     let device = Default::default();
///     burn::backend::wgpu::init_setup::<burn::backend::wgpu::graphics::Vulkan>(
///         &device,
///         Default::default(),
///     );
/// }
/// ```
/// will mean the given device (in this case the default) will be initialized to use Vulkan as the graphics API.
/// It's also possible to use an existing wgpu device, by using `init_device`.
///
/// # Notes
///
/// This version of the wgpu backend uses [burn_fusion] to compile and optimize streams of tensor
/// operations for improved performance.
///
/// You can disable the `fusion` feature flag to remove that functionality, which might be
/// necessary on `wasm` for now.
pub type Wgpu<F = f32, I = i32, B = u32> =
    burn_fusion::Fusion<CubeBackend<cubecl::wgpu::WgpuRuntime, F, I, B>>;
/// Tensor backend that uses the wgpu crate for executing GPU compute shaders.
///
/// This backend can target multiple graphics APIs, including:
///   - [Vulkan] on Linux, Windows, and Android.
///   - [OpenGL](crate::OpenGl) on Linux, Windows, and Android.
///   - [DirectX 12](crate::Dx12) on Windows.
///   - [Metal] on Apple hardware.
///   - [WebGPU](crate::WebGpu) on supported browsers and `wasm` runtimes.
///
/// To configure the wgpu backend, eg. to select what graphics API to use or what memory strategy to use,
/// you have to manually initialize the runtime. For example:
///
/// ```rust, ignore
/// fn custom_init() {
///     let device = Default::default();
///     burn::backend::wgpu::init_setup::<burn::backend::wgpu::graphics::Vulkan>(
///         &device,
///         Default::default(),
///     );
/// }
/// ```
/// will mean the given device (in this case the default) will be initialized to use Vulkan as the graphics API.
/// It's also possible to use an existing wgpu device, by using `init_device`.
///
/// # Notes
///
/// This version of the wgpu backend doesn't use [burn_fusion] to compile and optimize streams of tensor
/// operations.
///
/// You can enable the `fusion` feature flag to add that functionality, which might improve
/// performance.
pub type Wgpu<F = f32, I = i32, B = u32> = CubeBackend<cubecl::wgpu::WgpuRuntime, F, I, B>;
/// Tensor backend that leverages the Vulkan graphics API to execute GPU compute shaders compiled to SPIR-V.
pub type Vulkan<F = f32, I = i32, B = u8> = Wgpu<F, I, B>;
/// Tensor backend that uses the wgpu crate to execute GPU compute shaders written in WGSL.
pub type WebGpu<F = f32, I = i32, B = u32> = Wgpu<F, I, B>;
/// Tensor backend that leverages the Metal graphics API to execute GPU compute shaders compiled to MSL.
pub type Metal<F = f32, I = i32, B = u8> = Wgpu<F, I, B>;
    pub type TestRuntime = cubecl::wgpu::WgpuRuntime;
// Don't test `flex32` for now, burn sees it as `f32` but is actually `f16` precision, so it
// breaks a lot of tests from precision issues
