use alloc::{vec, vec::Vec};

use crate::Autodiff;
use crate::checkpoint::base::Checkpointer;
use crate::checkpoint::strategy::CheckpointStrategy;
use crate::grads::Gradients;
use crate::graph::NodeId;
use crate::ops::{Backward, Ops, unary};
use crate::tensor::AutodiffTensor;

use burn_backend::TensorMetadata;
use burn_backend::ops::attention::attention_fallback;
use burn_backend::ops::*;
use burn_backend::tensor::{FloatTensor, IntTensor};
use burn_backend::{Backend, get_device_settings};
use burn_std::{IntDType, Slice};

use super::OpsKind;

impl<B: Backend, C: CheckpointStrategy> ModuleOps<Autodiff<B, C>> for Autodiff<B, C> {
    fn embedding(weights: AutodiffTensor<B>, indices: IntTensor<B>) -> AutodiffTensor<B> {
        #[derive(Debug)]
        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,
            ) {
                let (weights, indices) = ops.state;

                unary::<B, _>(ops.parents, ops.node, grads, |grad| {
                    B::embedding_backward(weights, grad, indices)
                });
            }
        }

        match Embedding
            .prepare::<C>([weights.node])
            .compute_bound()
            .stateful()
        {
            OpsKind::Tracked(prep) => prep.finish(
                (weights.primitive.clone(), indices.clone()),
                B::embedding(weights.primitive, indices),
            ),
            OpsKind::UnTracked(prep) => prep.finish(B::embedding(weights.primitive, indices)),
        }
    }

    fn embedding_backward(
        _weights: AutodiffTensor<B>,
        _output: AutodiffTensor<B>,
        _indices: IntTensor<B>,
    ) -> AutodiffTensor<B> {
        panic!("Can't differentiate embedding backward.");
    }

    fn linear(
        x: AutodiffTensor<B>,
        weight: AutodiffTensor<B>,
        bias: Option<AutodiffTensor<B>>,
    ) -> AutodiffTensor<B> {
        #[derive(Debug)]
        struct LinearWithBias;
        #[derive(Debug)]
        struct LinearNoBias;

        impl<B: Backend> Backward<B, 3> for LinearWithBias {
            type State = (Option<NodeId>, Option<NodeId>);

            fn backward(
                self,
                ops: Ops<Self::State, 3>,
                grads: &mut Gradients,
                checkpointer: &mut Checkpointer,
            ) {
                let [node_x, node_weight, node_bias] = ops.parents;
                let grad = grads.consume::<B>(&ops.node);

                let (x_state, weight_state) = ops.state;
                let x = x_state
                    .map(|id| checkpointer.retrieve_node_output::<B::FloatTensorPrimitive>(id));
                let weight = weight_state
                    .map(|id| checkpointer.retrieve_node_output::<B::FloatTensorPrimitive>(id));

                if let Some(node) = node_x {
                    let grad = B::linear_x_backward(weight.unwrap(), grad.clone());
                    grads.register::<B>(node.id, grad)
                }
                if let Some(node) = node_weight {
                    let grad = B::linear_weight_backward(x.unwrap(), grad.clone());
                    grads.register::<B>(node.id, grad)
                }
                if let Some(node) = node_bias {
                    let grad = B::linear_bias_backward(grad);
                    grads.register::<B>(node.id, grad)
                }
            }
        }

        impl<B: Backend> Backward<B, 2> for LinearNoBias {
            type State = (Option<NodeId>, Option<NodeId>);

            fn backward(
                self,
                ops: Ops<Self::State, 2>,
                grads: &mut Gradients,
                checkpointer: &mut Checkpointer,
            ) {
                let [node_x, node_weight] = ops.parents;
                let grad = grads.consume::<B>(&ops.node);

                let (x_state, weight_state) = ops.state;
                let x = x_state
                    .map(|id| checkpointer.retrieve_node_output::<B::FloatTensorPrimitive>(id));
                let weight = weight_state
                    .map(|id| checkpointer.retrieve_node_output::<B::FloatTensorPrimitive>(id));

                if let Some(node) = node_x {
                    let grad = B::linear_x_backward(weight.unwrap(), grad.clone());
                    grads.register::<B>(node.id, grad)
                }
                if let Some(node) = node_weight {
                    let grad = B::linear_weight_backward(x.unwrap(), grad);
                    grads.register::<B>(node.id, grad)
                }
            }
        }

        let x_tracked = x.is_tracked();
        let weight_tracked = weight.is_tracked();

        match bias {
            Some(bias) => match LinearWithBias
                .prepare::<C>([x.node.clone(), weight.node.clone(), bias.node.clone()])
                .compute_bound()
                .stateful()
            {
                OpsKind::Tracked(mut prep) => {
                    // x is only needed to compute the weight gradient, and vice versa.
                    let x_state = weight_tracked.then(|| prep.checkpoint(&x));
                    let weight_state = x_tracked.then(|| prep.checkpoint(&weight));
                    prep.finish(
                        (x_state, weight_state),
                        B::linear(x.primitive, weight.primitive, Some(bias.primitive)),
                    )
                }
                OpsKind::UnTracked(prep) => prep.finish(B::linear(
                    x.primitive,
                    weight.primitive,
                    Some(bias.primitive),
                )),
            },
            None => match LinearNoBias
                .prepare::<C>([x.node.clone(), weight.node.clone()])
                .compute_bound()
                .stateful()
            {
                OpsKind::Tracked(mut prep) => {
                    let x_state = weight_tracked.then(|| prep.checkpoint(&x));
                    let weight_state = x_tracked.then(|| prep.checkpoint(&weight));
                    prep.finish(
                        (x_state, weight_state),
                        B::linear(x.primitive, weight.primitive, None),
                    )
                }
                OpsKind::UnTracked(prep) => {
                    prep.finish(B::linear(x.primitive, weight.primitive, None))
                }
            },
        }
    }

    fn linear_x_backward(
        _weight: AutodiffTensor<B>,
        _output_grad: AutodiffTensor<B>,
    ) -> AutodiffTensor<B> {
        panic!("Can't differentiate linear_x_backward.");
    }

    fn linear_weight_backward(
        _x: AutodiffTensor<B>,
        _output_grad: AutodiffTensor<B>,
    ) -> AutodiffTensor<B> {
        panic!("Can't differentiate linear_weight_backward.");
    }

    fn linear_bias_backward(_output_grad: AutodiffTensor<B>) -> AutodiffTensor<B> {
        panic!("Can't differentiate linear_bias_backward.");
    }

    fn conv1d(
        x: AutodiffTensor<B>,
        weight: AutodiffTensor<B>,
        bias: Option<AutodiffTensor<B>>,
        options: ConvOptions<1>,
    ) -> AutodiffTensor<B> {
        #[derive(Debug)]
        struct Conv1DWithBias;
        #[derive(Debug)]
        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,
            ) {
                let [node_x, node_weight, node_bias] = ops.parents;
                let grad = grads.consume::<B>(&ops.node);

                let (x_state, weight_state, bias_state, options) = ops.state;
                let x = checkpointer.retrieve_node_output::<B::FloatTensorPrimitive>(x_state);
                let weight =
                    checkpointer.retrieve_node_output::<B::FloatTensorPrimitive>(weight_state);
                let bias = checkpointer.retrieve_node_output::<B::FloatTensorPrimitive>(bias_state);

                if let Some(node) = node_x {
                    let grad = B::conv1d_x_backward(
                        x.clone(),
                        weight.clone(),
                        grad.clone(),
                        options.clone(),
                    );
                    grads.register::<B>(node.id, grad)
                }
                if let Some(node) = node_weight {
                    let grad = B::conv1d_weight_backward(x.clone(), weight, grad.clone(), options);
                    grads.register::<B>(node.id, grad)
                }
                if let Some(node) = node_bias {
                    let grad = B::conv1d_bias_backward(x, bias, grad);
                    grads.register::<B>(node.id, grad)
                }
            }
        }

        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,
            ) {
                let [node_x, node_weight] = ops.parents;
                let grad = grads.consume::<B>(&ops.node);

                let (x_state, weight_state, options) = ops.state;
                let x = checkpointer.retrieve_node_output::<B::FloatTensorPrimitive>(x_state);
                let weight =
                    checkpointer.retrieve_node_output::<B::FloatTensorPrimitive>(weight_state);

                if let Some(node) = node_x {
                    let grad = B::conv1d_x_backward(
                        x.clone(),
                        weight.clone(),
                        grad.clone(),
                        options.clone(),
                    );
                    grads.register::<B>(node.id, grad)
                }
                if let Some(node) = node_weight {
                    let grad = B::conv1d_weight_backward(x, weight, grad, options);
                    grads.register::<B>(node.id, grad)
                }
            }
        }
        match bias {
            Some(bias) => match Conv1DWithBias
                .prepare::<C>([x.node.clone(), weight.node.clone(), bias.node.clone()])
                .compute_bound()
                .stateful()
            {
                OpsKind::Tracked(mut prep) => {
                    let x_state = prep.checkpoint(&x);
                    let weight_state = prep.checkpoint(&weight);
                    let bias_state = prep.checkpoint(&bias);
                    prep.finish(
                        (x_state, weight_state, bias_state, options.clone()),
                        B::conv1d(x.primitive, weight.primitive, Some(bias.primitive), options),
                    )
                }
                OpsKind::UnTracked(prep) => prep.finish(B::conv1d(
                    x.primitive,
                    weight.primitive,
                    Some(bias.primitive),
                    options,
                )),
            },
            None => match Conv1DNoBias
                .prepare::<C>([x.node.clone(), weight.node.clone()])
                .compute_bound()
                .stateful()
            {
                OpsKind::Tracked(mut prep) => {
                    let x_state = prep.checkpoint(&x);
                    let weight_state = prep.checkpoint(&weight);
                    prep.finish(
                        (x_state, weight_state, options.clone()),
                        B::conv1d(x.primitive, weight.primitive, None, options),
                    )
                }
                OpsKind::UnTracked(prep) => {
                    prep.finish(B::conv1d(x.primitive, weight.primitive, None, options))
                }
            },
        }
    }

    fn conv_transpose1d(
        x: AutodiffTensor<B>,
        weight: AutodiffTensor<B>,
        bias: Option<AutodiffTensor<B>>,
        options: ConvTransposeOptions<1>,
    ) -> AutodiffTensor<B> {
        #[derive(Debug)]
        struct ConvTranspose1DWithBias;
        #[derive(Debug)]
        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,
            ) {
                let [node_x, node_weight, node_bias] = ops.parents;
                let grad = grads.consume::<B>(&ops.node);

                let (x_state, weight_state, bias_state, options) = ops.state;
                let x = checkpointer.retrieve_node_output::<B::FloatTensorPrimitive>(x_state);
                let weight =
                    checkpointer.retrieve_node_output::<B::FloatTensorPrimitive>(weight_state);
                let bias = checkpointer.retrieve_node_output::<B::FloatTensorPrimitive>(bias_state);

                if let Some(node) = node_x {
                    let grad = B::conv_transpose1d_x_backward(
                        weight.clone(),
                        grad.clone(),
                        options.clone(),
                    );
                    grads.register::<B>(node.id, grad)
                }
                if let Some(node) = node_weight {
                    let grad = B::conv_transpose1d_weight_backward(
                        x.clone(),
                        weight,
                        grad.clone(),
                        options,
                    );
                    grads.register::<B>(node.id, grad)
                }
                if let Some(node) = node_bias {
                    let grad = B::conv_transpose1d_bias_backward(x, bias, grad);
                    grads.register::<B>(node.id, grad)
                }
            }
        }

        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,
            ) {
                let [node_x, node_weight] = ops.parents;
                let grad = grads.consume::<B>(&ops.node);

                let (x_state, weight_state, options) = ops.state;
                let x = checkpointer.retrieve_node_output::<B::FloatTensorPrimitive>(x_state);
                let weight =
                    checkpointer.retrieve_node_output::<B::FloatTensorPrimitive>(weight_state);

                if let Some(node) = node_x {
                    let grad = B::conv_transpose1d_x_backward(
                        weight.clone(),
                        grad.clone(),
                        options.clone(),
                    );
                    grads.register::<B>(node.id, grad)
                }
                if let Some(node) = node_weight {
                    let grad = B::conv_transpose1d_weight_backward(x, weight, grad, options);
                    grads.register::<B>(node.id, grad)
                }
            }
        }

        match bias {
            Some(bias) => match ConvTranspose1DWithBias
                .prepare::<C>([x.node.clone(), weight.node.clone(), bias.node.clone()])
                .compute_bound()
                .stateful()
            {
                OpsKind::Tracked(mut prep) => {
                    let x_state = prep.checkpoint(&x);
                    let weight_state = prep.checkpoint(&weight);
                    let bias_state = prep.checkpoint(&bias);
                    prep.finish(
                        (x_state, weight_state, bias_state, options.clone()),
                        B::conv_transpose1d(
                            x.primitive,
                            weight.primitive,
                            Some(bias.primitive),
                            options,
                        ),
                    )
                }
                OpsKind::UnTracked(prep) => prep.finish(B::conv_transpose1d(
                    x.primitive,
                    weight.primitive,
                    Some(bias.primitive),
                    options,
                )),
            },
            None => match ConvTranspose1DNoBias
                .prepare::<C>([x.node.clone(), weight.node.clone()])
                .compute_bound()
                .stateful()
            {
                OpsKind::Tracked(mut prep) => {
                    let x_state = prep.checkpoint(&x);
                    let weight_state = prep.checkpoint(&weight);
                    prep.finish(
                        (x_state, weight_state, options.clone()),
                        B::conv_transpose1d(x.primitive, weight.primitive, None, options),
                    )
                }
                OpsKind::UnTracked(prep) => prep.finish(B::conv_transpose1d(
                    x.primitive,
                    weight.primitive,
                    None,
                    options,
                )),
            },
        }
    }

    fn conv2d(
        x: AutodiffTensor<B>,
        weight: AutodiffTensor<B>,
        bias: Option<AutodiffTensor<B>>,
        options: ConvOptions<2>,
    ) -> AutodiffTensor<B> {
        #[derive(Debug)]
        struct Conv2DWithBias;
        #[derive(Debug)]
        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,
            ) {
                let [node_x, node_weight, node_bias] = ops.parents;
                let grad = grads.consume::<B>(&ops.node);

                let (x_state, weight_state, bias_state, options) = ops.state;
                let x = checkpointer.retrieve_node_output::<B::FloatTensorPrimitive>(x_state);
                let weight =
                    checkpointer.retrieve_node_output::<B::FloatTensorPrimitive>(weight_state);
                let bias = checkpointer.retrieve_node_output::<B::FloatTensorPrimitive>(bias_state);

                if let Some(node) = node_x {
                    let grad = B::conv2d_x_backward(
                        x.clone(),
                        weight.clone(),
                        grad.clone(),
                        options.clone(),
                    );
                    grads.register::<B>(node.id, grad)
                }
                if let Some(node) = node_weight {
                    let grad =
                        B::conv2d_weight_backward(x.clone(), weight.clone(), grad.clone(), options);
                    grads.register::<B>(node.id, grad)
                }
                if let Some(node) = node_bias {
                    let grad = B::conv2d_bias_backward(x, bias, grad);
                    grads.register::<B>(node.id, grad)
                }
            }
        }

        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,
            ) {
                let [node_x, node_weight] = ops.parents;
                let grad = grads.consume::<B>(&ops.node);

                let (x_state, weight_state, options) = ops.state;
                let x = checkpointer.retrieve_node_output::<B::FloatTensorPrimitive>(x_state);
                let weight =
                    checkpointer.retrieve_node_output::<B::FloatTensorPrimitive>(weight_state);

                if let Some(node) = node_x {
                    let grad = B::conv2d_x_backward(
                        x.clone(),
                        weight.clone(),
                        grad.clone(),
                        options.clone(),
                    );
                    grads.register::<B>(node.id, grad)
                }
                if let Some(node) = node_weight {
                    let grad = B::conv2d_weight_backward(x, weight, grad, options);
                    grads.register::<B>(node.id, grad)
                }
            }
        }

        match bias {
            Some(bias) => match Conv2DWithBias
                .prepare::<C>([x.node.clone(), weight.node.clone(), bias.node.clone()])
                .compute_bound()
                .stateful()
            {
                OpsKind::Tracked(mut prep) => {
                    let x_state = prep.checkpoint(&x);
                    let weight_state = prep.checkpoint(&weight);
                    let bias_state = prep.checkpoint(&bias);
                    prep.finish(
                        (x_state, weight_state, bias_state, options.clone()),
                        B::conv2d(x.primitive, weight.primitive, Some(bias.primitive), options),
                    )
                }
                OpsKind::UnTracked(prep) => prep.finish(B::conv2d(
                    x.primitive,
                    weight.primitive,
                    Some(bias.primitive),
                    options,
                )),
            },
            None => match Conv2DNoBias
                .prepare::<C>([x.node.clone(), weight.node.clone()])
                .compute_bound()
                .stateful()
            {
                OpsKind::Tracked(mut prep) => {
                    let x_state = prep.checkpoint(&x);
                    let weight_state = prep.checkpoint(&weight);
                    prep.finish(
                        (x_state, weight_state, options.clone()),
                        B::conv2d(x.primitive, weight.primitive, None, options),
                    )
                }

                OpsKind::UnTracked(prep) => {
                    prep.finish(B::conv2d(x.primitive, weight.primitive, None, options))
                }
            },
        }
    }

    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> {
        #[derive(Debug)]
        struct DeformConv2DWithMaskWithBias;
        #[derive(Debug)]
        struct DeformConv2DWithMaskNoBias;
        #[derive(Debug)]
        struct DeformConv2DNoMaskWithBias;
        #[derive(Debug)]
        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,
            ) {
                let [node_x, node_offset, node_weight, node_mask, node_bias] = ops.parents;
                let grad = grads.consume::<B>(&ops.node);

                let (x_state, offset_state, weight_state, mask_state, bias_state, options) =
                    ops.state;
                let x = checkpointer.retrieve_node_output(x_state);
                let offset = checkpointer.retrieve_node_output(offset_state);
                let weight = checkpointer.retrieve_node_output(weight_state);
                let mask = Some(checkpointer.retrieve_node_output(mask_state));
                let bias = Some(checkpointer.retrieve_node_output(bias_state));

                let backward =
                    B::deform_conv2d_backward(x, offset, weight, mask, bias, grad, options);

                if let Some(node) = node_x {
                    grads.register::<B>(node.id, backward.x_grad)
                }
                if let Some(node) = node_offset {
                    grads.register::<B>(node.id, backward.offset_grad)
                }
                if let Some(node) = node_weight {
                    grads.register::<B>(node.id, backward.weight_grad)
                }
                if let Some(node) = node_mask {
                    grads.register::<B>(node.id, backward.mask_grad.unwrap())
                }
                if let Some(node) = node_bias {
                    grads.register::<B>(node.id, backward.bias_grad.unwrap())
                }
            }
        }

        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,
            ) {
                let [node_x, node_offset, node_weight, node_mask] = ops.parents;
                let grad = grads.consume::<B>(&ops.node);

                let (x_state, offset_state, weight_state, mask_state, options) = ops.state;
                let x = checkpointer.retrieve_node_output(x_state);
                let offset = checkpointer.retrieve_node_output(offset_state);
                let weight = checkpointer.retrieve_node_output(weight_state);
                let mask = Some(checkpointer.retrieve_node_output(mask_state));

                let backward =
                    B::deform_conv2d_backward(x, offset, weight, mask, None, grad, options);

                if let Some(node) = node_x {
                    grads.register::<B>(node.id, backward.x_grad)
                }
                if let Some(node) = node_offset {
                    grads.register::<B>(node.id, backward.offset_grad)
                }
                if let Some(node) = node_weight {
                    grads.register::<B>(node.id, backward.weight_grad)
                }
                if let Some(node) = node_mask {
                    grads.register::<B>(node.id, backward.mask_grad.unwrap())
                }
            }
        }

        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,
            ) {
                let [node_x, node_offset, node_weight, node_bias] = ops.parents;
                let grad = grads.consume::<B>(&ops.node);

                let (x_state, offset_state, weight_state, bias_state, options) = ops.state;
                let x = checkpointer.retrieve_node_output(x_state);
                let offset = checkpointer.retrieve_node_output(offset_state);
                let weight = checkpointer.retrieve_node_output(weight_state);
                let bias = Some(checkpointer.retrieve_node_output(bias_state));

                let backward =
                    B::deform_conv2d_backward(x, offset, weight, None, bias, grad, options);

                if let Some(node) = node_x {
                    grads.register::<B>(node.id, backward.x_grad)
                }
                if let Some(node) = node_offset {
                    grads.register::<B>(node.id, backward.offset_grad)
                }
                if let Some(node) = node_weight {
                    grads.register::<B>(node.id, backward.weight_grad)
                }
                if let Some(node) = node_bias {
                    grads.register::<B>(node.id, backward.bias_grad.unwrap())
                }
            }
        }

        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,
            ) {
                let [node_x, node_offset, node_weight] = ops.parents;
                let grad = grads.consume::<B>(&ops.node);

                let (x_state, offset_state, weight_state, options) = ops.state;
                let x = checkpointer.retrieve_node_output(x_state);
                let offset = checkpointer.retrieve_node_output(offset_state);
                let weight = checkpointer.retrieve_node_output(weight_state);

                let backward =
                    B::deform_conv2d_backward(x, offset, weight, None, None, grad, options);

                if let Some(node) = node_x {
                    grads.register::<B>(node.id, backward.x_grad)
                }
                if let Some(node) = node_offset {
                    grads.register::<B>(node.id, backward.offset_grad)
                }
                if let Some(node) = node_weight {
                    grads.register::<B>(node.id, backward.weight_grad)
                }
            }
        }

        match (mask, bias) {
            (Some(mask), Some(bias)) => match DeformConv2DWithMaskWithBias
                .prepare::<C>([
                    x.node.clone(),
                    offset.node.clone(),
                    weight.node.clone(),
                    mask.node.clone(),
                    bias.node.clone(),
                ])
                .compute_bound()
                .stateful()
            {
                OpsKind::Tracked(mut prep) => {
                    let x_state = prep.checkpoint(&x);
                    let offset_state = prep.checkpoint(&offset);
                    let weight_state = prep.checkpoint(&weight);
                    let mask_state = prep.checkpoint(&mask);
                    let bias_state = prep.checkpoint(&bias);
                    prep.finish(
                        (
                            x_state,
                            offset_state,
                            weight_state,
                            mask_state,
                            bias_state,
                            options.clone(),
                        ),
                        B::deform_conv2d(
                            x.primitive,
                            offset.primitive,
                            weight.primitive,
                            Some(mask.primitive),
                            Some(bias.primitive),
                            options,
                        ),
                    )
                }
                OpsKind::UnTracked(prep) => prep.finish(B::deform_conv2d(
                    x.primitive,
                    offset.primitive,
                    weight.primitive,
                    Some(mask.primitive),
                    Some(bias.primitive),
                    options,
                )),
            },
            (Some(mask), None) => match DeformConv2DWithMaskNoBias
                .prepare::<C>([
                    x.node.clone(),
                    offset.node.clone(),
                    weight.node.clone(),
                    mask.node.clone(),
                ])
                .compute_bound()
                .stateful()
            {
                OpsKind::Tracked(mut prep) => {
                    let x_state = prep.checkpoint(&x);
                    let offset_state = prep.checkpoint(&offset);
                    let weight_state = prep.checkpoint(&weight);
                    let mask_state = prep.checkpoint(&mask);
                    prep.finish(
                        (
                            x_state,
                            offset_state,
                            weight_state,
                            mask_state,
                            options.clone(),
                        ),
                        B::deform_conv2d(
                            x.primitive,
                            offset.primitive,
                            weight.primitive,
                            Some(mask.primitive),
                            None,
                            options,
                        ),
                    )
                }
                OpsKind::UnTracked(prep) => prep.finish(B::deform_conv2d(
                    x.primitive,
                    offset.primitive,
                    weight.primitive,
                    Some(mask.primitive),
                    None,
                    options,
                )),
            },
            (None, Some(bias)) => match DeformConv2DNoMaskWithBias
                .prepare::<C>([
                    x.node.clone(),
                    offset.node.clone(),
                    weight.node.clone(),
                    bias.node.clone(),
                ])
                .compute_bound()
                .stateful()
            {
                OpsKind::Tracked(mut prep) => {
                    let x_state = prep.checkpoint(&x);
                    let offset_state = prep.checkpoint(&offset);
                    let weight_state = prep.checkpoint(&weight);
                    let bias_state = prep.checkpoint(&bias);
                    prep.finish(
                        (
                            x_state,
                            offset_state,
                            weight_state,
                            bias_state,
                            options.clone(),
                        ),
                        B::deform_conv2d(
                            x.primitive,
                            offset.primitive,
                            weight.primitive,
                            None,
                            Some(bias.primitive),
                            options,
                        ),
                    )
                }
                OpsKind::UnTracked(prep) => prep.finish(B::deform_conv2d(
                    x.primitive,
                    offset.primitive,
                    weight.primitive,
                    None,
                    Some(bias.primitive),
                    options,
                )),
            },
            (None, None) => match DeformConv2DNoMaskNoBias
                .prepare::<C>([x.node.clone(), offset.node.clone(), weight.node.clone()])
                .compute_bound()
                .stateful()
            {
                OpsKind::Tracked(mut prep) => {
                    let x_state = prep.checkpoint(&x);
                    let offset_state = prep.checkpoint(&offset);
                    let weight_state = prep.checkpoint(&weight);
                    prep.finish(
                        (x_state, offset_state, weight_state, options.clone()),
                        B::deform_conv2d(
                            x.primitive,
                            offset.primitive,
                            weight.primitive,
                            None,
                            None,
                            options,
                        ),
                    )
                }
                OpsKind::UnTracked(prep) => prep.finish(B::deform_conv2d(
                    x.primitive,
                    offset.primitive,
                    weight.primitive,
                    None,
                    None,
                    options,
                )),
            },
        }
    }

    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> {
        panic!("Can't differentiate deform conv 2d backward.");
    }

    fn conv_transpose2d(
        x: AutodiffTensor<B>,
        weight: AutodiffTensor<B>,
        bias: Option<AutodiffTensor<B>>,
        options: ConvTransposeOptions<2>,
    ) -> AutodiffTensor<B> {
        #[derive(Debug)]
        struct ConvTranspose2DWithBias;
        #[derive(Debug)]
        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,
            ) {
                let [node_x, node_weight, node_bias] = ops.parents;
                let grad = grads.consume::<B>(&ops.node);

                let (x_state, weight_state, bias_state, options) = ops.state;
                let x = checkpointer.retrieve_node_output::<B::FloatTensorPrimitive>(x_state);
                let weight =
                    checkpointer.retrieve_node_output::<B::FloatTensorPrimitive>(weight_state);
                let bias = checkpointer.retrieve_node_output::<B::FloatTensorPrimitive>(bias_state);

                if let Some(node) = node_x {
                    let grad = B::conv_transpose2d_x_backward(
                        weight.clone(),
                        grad.clone(),
                        options.clone(),
                    );
                    grads.register::<B>(node.id, grad)
                }
                if let Some(node) = node_weight {
                    let grad = B::conv_transpose2d_weight_backward(
                        x.clone(),
                        weight,
                        grad.clone(),
                        options,
                    );
                    grads.register::<B>(node.id, grad)
                }
                if let Some(node) = node_bias {
                    let grad = B::conv_transpose2d_bias_backward(x, bias, grad);
                    grads.register::<B>(node.id, grad)
                }
            }
        }

        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,
            ) {
                let [node_x, node_weight] = ops.parents;
                let grad = grads.consume::<B>(&ops.node);

                let (x_state, weight_state, options) = ops.state;
                let x = checkpointer.retrieve_node_output::<B::FloatTensorPrimitive>(x_state);
                let weight =
                    checkpointer.retrieve_node_output::<B::FloatTensorPrimitive>(weight_state);

                if let Some(node) = node_x {
                    let grad = B::conv_transpose2d_x_backward(
                        weight.clone(),
                        grad.clone(),
                        options.clone(),
                    );
                    grads.register::<B>(node.id, grad)
                }
                if let Some(node) = node_weight {
                    let grad = B::conv_transpose2d_weight_backward(x, weight, grad, options);
                    grads.register::<B>(node.id, grad)
                }
            }
        }

        match bias {
            Some(bias) => match ConvTranspose2DWithBias
                .prepare::<C>([x.node.clone(), weight.node.clone(), bias.node.clone()])
                .compute_bound()
                .stateful()
            {
                OpsKind::Tracked(mut prep) => {
                    let x_state = prep.checkpoint(&x);
                    let weight_state = prep.checkpoint(&weight);
                    let bias_state = prep.checkpoint(&bias);

                    prep.finish(
                        (x_state, weight_state, bias_state, options.clone()),
                        B::conv_transpose2d(
                            x.primitive,
                            weight.primitive,
                            Some(bias.primitive),
                            options,
                        ),
                    )
                }
                OpsKind::UnTracked(prep) => prep.finish(B::conv_transpose2d(
                    x.primitive,
                    weight.primitive,
                    Some(bias.primitive),
                    options,
                )),
            },
            None => match ConvTranspose2DNoBias
                .prepare::<C>([x.node.clone(), weight.node.clone()])
                .compute_bound()
                .stateful()
            {
                OpsKind::Tracked(mut prep) => {
                    let x_state = prep.checkpoint(&x);
                    let weight_state = prep.checkpoint(&weight);

                    prep.finish(
                        (x_state, weight_state, options.clone()),
                        B::conv_transpose2d(x.primitive, weight.primitive, None, options),
                    )
                }
                OpsKind::UnTracked(prep) => prep.finish(B::conv_transpose2d(
                    x.primitive,
                    weight.primitive,
                    None,
                    options,
                )),
            },
        }
    }

    fn conv3d(
        x: AutodiffTensor<B>,
        weight: AutodiffTensor<B>,
        bias: Option<AutodiffTensor<B>>,
        options: ConvOptions<3>,
    ) -> AutodiffTensor<B> {
        #[derive(Debug)]
        struct Conv3DWithBias;
        #[derive(Debug)]
        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,
            ) {
                let [node_x, node_weight, node_bias] = ops.parents;
                let grad = grads.consume::<B>(&ops.node);

                let (x_state, weight_state, bias_state, options) = ops.state;
                let x = checkpointer.retrieve_node_output::<B::FloatTensorPrimitive>(x_state);
                let weight =
                    checkpointer.retrieve_node_output::<B::FloatTensorPrimitive>(weight_state);
                let bias = checkpointer.retrieve_node_output::<B::FloatTensorPrimitive>(bias_state);

                if let Some(node) = node_x {
                    let grad = B::conv3d_x_backward(
                        x.clone(),
                        weight.clone(),
                        grad.clone(),
                        options.clone(),
                    );
                    grads.register::<B>(node.id, grad)
                }
                if let Some(node) = node_weight {
                    let grad =
                        B::conv3d_weight_backward(x.clone(), weight.clone(), grad.clone(), options);
                    grads.register::<B>(node.id, grad)
                }
                if let Some(node) = node_bias {
                    let grad = B::conv3d_bias_backward(x, bias, grad);
                    grads.register::<B>(node.id, grad)
                }
            }
        }

        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,
            ) {
                let [node_x, node_weight] = ops.parents;
                let grad = grads.consume::<B>(&ops.node);

                let (x_state, weight_state, options) = ops.state;
                let x = checkpointer.retrieve_node_output::<B::FloatTensorPrimitive>(x_state);
                let weight =
                    checkpointer.retrieve_node_output::<B::FloatTensorPrimitive>(weight_state);

                if let Some(node) = node_x {
                    let grad = B::conv3d_x_backward(
                        x.clone(),
                        weight.clone(),
                        grad.clone(),
                        options.clone(),
                    );
                    grads.register::<B>(node.id, grad)
                }
                if let Some(node) = node_weight {
                    let grad = B::conv3d_weight_backward(x, weight, grad, options);
                    grads.register::<B>(node.id, grad)
                }
            }
        }

        match bias {
            Some(bias) => match Conv3DWithBias
                .prepare::<C>([x.node.clone(), weight.node.clone(), bias.node.clone()])
                .compute_bound()
                .stateful()
            {
                OpsKind::Tracked(mut prep) => {
                    let x_state = prep.checkpoint(&x);
                    let weight_state = prep.checkpoint(&weight);
                    let bias_state = prep.checkpoint(&bias);
                    prep.finish(
                        (x_state, weight_state, bias_state, options.clone()),
                        B::conv3d(x.primitive, weight.primitive, Some(bias.primitive), options),
                    )
                }
                OpsKind::UnTracked(prep) => prep.finish(B::conv3d(
                    x.primitive,
                    weight.primitive,
                    Some(bias.primitive),
                    options,
                )),
            },
            None => match Conv3DNoBias
                .prepare::<C>([x.node.clone(), weight.node.clone()])
                .compute_bound()
                .stateful()
            {
                OpsKind::Tracked(mut prep) => {
                    let x_state = prep.checkpoint(&x);
                    let weight_state = prep.checkpoint(&weight);
                    prep.finish(
                        (x_state, weight_state, options.clone()),
                        B::conv3d(x.primitive, weight.primitive, None, options),
                    )
                }

                OpsKind::UnTracked(prep) => {
                    prep.finish(B::conv3d(x.primitive, weight.primitive, None, options))
                }
            },
        }
    }

    fn conv_transpose3d(
        x: AutodiffTensor<B>,
        weight: AutodiffTensor<B>,
        bias: Option<AutodiffTensor<B>>,
        options: ConvTransposeOptions<3>,
    ) -> AutodiffTensor<B> {
        #[derive(Debug)]
        struct ConvTranspose3DWithBias;
        #[derive(Debug)]
        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,
            ) {
                let [node_x, node_weight, node_bias] = ops.parents;
                let grad = grads.consume::<B>(&ops.node);

                let (x_state, weight_state, bias_state, options) = ops.state;
                let x = checkpointer.retrieve_node_output::<B::FloatTensorPrimitive>(x_state);
                let weight =
                    checkpointer.retrieve_node_output::<B::FloatTensorPrimitive>(weight_state);
                let bias = checkpointer.retrieve_node_output::<B::FloatTensorPrimitive>(bias_state);

                if let Some(node) = node_x {
                    let grad = B::conv_transpose3d_x_backward(
                        weight.clone(),
                        grad.clone(),
                        options.clone(),
                    );
                    grads.register::<B>(node.id, grad)
                }
                if let Some(node) = node_weight {
                    let grad = B::conv_transpose3d_weight_backward(
                        x.clone(),
                        weight,
                        grad.clone(),
                        options,
                    );
                    grads.register::<B>(node.id, grad)
                }
                if let Some(node) = node_bias {
                    let grad = B::conv_transpose3d_bias_backward(x, bias, grad);
                    grads.register::<B>(node.id, grad)
                }
            }
        }

        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,
            ) {
                let [node_x, node_weight] = ops.parents;
                let grad = grads.consume::<B>(&ops.node);

                let (x_state, weight_state, options) = ops.state;
                let x = checkpointer.retrieve_node_output::<B::FloatTensorPrimitive>(x_state);
                let weight =
                    checkpointer.retrieve_node_output::<B::FloatTensorPrimitive>(weight_state);

                if let Some(node) = node_x {
                    let grad = B::conv_transpose3d_x_backward(
                        weight.clone(),
                        grad.clone(),
                        options.clone(),
                    );
                    grads.register::<B>(node.id, grad)
                }
                if let Some(node) = node_weight {
                    let grad = B::conv_transpose3d_weight_backward(x, weight, grad, options);
                    grads.register::<B>(node.id, grad)
                }
            }
        }

        match bias {
            Some(bias) => match ConvTranspose3DWithBias
                .prepare::<C>([x.node.clone(), weight.node.clone(), bias.node.clone()])
                .compute_bound()
                .stateful()
            {
                OpsKind::Tracked(mut prep) => {
                    let x_state = prep.checkpoint(&x);
                    let weight_state = prep.checkpoint(&weight);
                    let bias_state = prep.checkpoint(&bias);

                    prep.finish(
                        (x_state, weight_state, bias_state, options.clone()),
                        B::conv_transpose3d(
                            x.primitive,
                            weight.primitive,
                            Some(bias.primitive),
                            options,
                        ),
                    )
                }
                OpsKind::UnTracked(prep) => prep.finish(B::conv_transpose3d(
                    x.primitive,
                    weight.primitive,
                    Some(bias.primitive),
                    options,
                )),
            },
            None => match ConvTranspose3DNoBias
                .prepare::<C>([x.node.clone(), weight.node.clone()])
                .compute_bound()
                .stateful()
            {
                OpsKind::Tracked(mut prep) => {
                    let x_state = prep.checkpoint(&x);
                    let weight_state = prep.checkpoint(&weight);

                    prep.finish(
                        (x_state, weight_state, options.clone()),
                        B::conv_transpose3d(x.primitive, weight.primitive, None, options),
                    )
                }
                OpsKind::UnTracked(prep) => prep.finish(B::conv_transpose3d(
                    x.primitive,
                    weight.primitive,
                    None,
                    options,
                )),
            },
        }
    }

    // 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,
        ceil_mode: bool,
    ) -> AutodiffTensor<B> {
        #[derive(Debug)]
        struct AvgPool1D;

        impl<B: Backend> Backward<B, 1> for AvgPool1D {
            type State = (NodeId, usize, usize, usize, bool, bool);

            fn backward(
                self,
                ops: Ops<Self::State, 1>,
                grads: &mut Gradients,
                checkpointer: &mut Checkpointer,
            ) {
                let [node_parent] = ops.parents;
                let grad = grads.consume::<B>(&ops.node);
                let (x_state, kernel_size, stride, padding, count_include_pad, ceil_mode) =
                    ops.state;
                let x = checkpointer.retrieve_node_output(x_state);

                if let Some(node) = node_parent {
                    let grad = B::avg_pool1d_backward(
                        x,
                        grad,
                        kernel_size,
                        stride,
                        padding,
                        count_include_pad,
                        ceil_mode,
                    );
                    grads.register::<B>(node.id, grad);
                }
            }
        }

        match AvgPool1D
            .prepare::<C>([x.node.clone()])
            .compute_bound()
            .stateful()
        {
            OpsKind::Tracked(mut prep) => {
                let x_state = prep.checkpoint(&x);
                prep.finish(
                    (
                        x_state,
                        kernel_size,
                        stride,
                        padding,
                        count_include_pad,
                        ceil_mode,
                    ),
                    B::avg_pool1d(
                        x.primitive.clone(),
                        kernel_size,
                        stride,
                        padding,
                        count_include_pad,
                        ceil_mode,
                    ),
                )
            }
            OpsKind::UnTracked(prep) => prep.finish(B::avg_pool1d(
                x.primitive,
                kernel_size,
                stride,
                padding,
                count_include_pad,
                ceil_mode,
            )),
        }
    }

    fn avg_pool2d(
        x: AutodiffTensor<B>,
        kernel_size: [usize; 2],
        stride: [usize; 2],
        padding: [usize; 2],
        count_include_pad: bool,
        ceil_mode: bool,
    ) -> AutodiffTensor<B> {
        #[derive(Debug)]
        struct AvgPool2D;

        impl<B: Backend> Backward<B, 1> for AvgPool2D {
            type State = (NodeId, [usize; 2], [usize; 2], [usize; 2], bool, bool);

            fn backward(
                self,
                ops: Ops<Self::State, 1>,
                grads: &mut Gradients,
                checkpointer: &mut Checkpointer,
            ) {
                let [node_parent] = ops.parents;
                let grad = grads.consume::<B>(&ops.node);
                let (x_state, kernel_size, stride, padding, count_include_pad, ceil_mode) =
                    ops.state;
                let x = checkpointer.retrieve_node_output(x_state);

                if let Some(node) = node_parent {
                    let grad = B::avg_pool2d_backward(
                        x,
                        grad,
                        kernel_size,
                        stride,
                        padding,
                        count_include_pad,
                        ceil_mode,
                    );
                    grads.register::<B>(node.id, grad);
                }
            }
        }

        match AvgPool2D
            .prepare::<C>([x.node.clone()])
            .compute_bound()
            .stateful()
        {
            OpsKind::Tracked(mut prep) => {
                let x_state = prep.checkpoint(&x);
                prep.finish(
                    (
                        x_state,
                        kernel_size,
                        stride,
                        padding,
                        count_include_pad,
                        ceil_mode,
                    ),
                    B::avg_pool2d(
                        x.primitive.clone(),
                        kernel_size,
                        stride,
                        padding,
                        count_include_pad,
                        ceil_mode,
                    ),
                )
            }
            OpsKind::UnTracked(prep) => prep.finish(B::avg_pool2d(
                x.primitive,
                kernel_size,
                stride,
                padding,
                count_include_pad,
                ceil_mode,
            )),
        }
    }

    fn avg_pool2d_backward(
        _x: AutodiffTensor<B>,
        _grad: AutodiffTensor<B>,
        _kernel_size: [usize; 2],
        _stride: [usize; 2],
        _padding: [usize; 2],
        _count_include_pad: bool,
        _ceil_mode: bool,
    ) -> AutodiffTensor<B> {
        panic!("Can't differentiate avg pool 2d backward.");
    }

    fn max_pool1d(
        x: AutodiffTensor<B>,
        kernel_size: usize,
        stride: usize,
        padding: usize,
        dilation: usize,
        ceil_mode: bool,
    ) -> AutodiffTensor<B> {
        match MaxPool1D
            .prepare::<C>([x.node.clone()])
            .compute_bound()
            .stateful()
        {
            OpsKind::Tracked(mut prep) => {
                let x_state = prep.checkpoint(&x);
                let settings = get_device_settings::<B>(&B::float_device(&x.primitive));
                let output = B::max_pool1d_with_indices(
                    x.primitive,
                    kernel_size,
                    stride,
                    padding,
                    dilation,
                    ceil_mode,
                    settings.int_dtype,
                );
                prep.finish(
                    (
                        x_state,
                        output.indices,
                        kernel_size,
                        stride,
                        padding,
                        dilation,
                        ceil_mode,
                    ),
                    output.output,
                )
            }
            OpsKind::UnTracked(prep) => prep.finish(B::max_pool1d(
                x.primitive,
                kernel_size,
                stride,
                padding,
                dilation,
                ceil_mode,
            )),
        }
    }

    fn max_pool1d_with_indices(
        x: AutodiffTensor<B>,
        kernel_size: usize,
        stride: usize,
        padding: usize,
        dilation: usize,
        ceil_mode: bool,
        int_dtype: IntDType,
    ) -> MaxPool1dWithIndices<Self> {
        match MaxPool1D
            .prepare::<C>([x.node.clone()])
            .compute_bound()
            .stateful()
        {
            OpsKind::Tracked(mut prep) => {
                let x_state = prep.checkpoint(&x);
                let output = B::max_pool1d_with_indices(
                    x.primitive,
                    kernel_size,
                    stride,
                    padding,
                    dilation,
                    ceil_mode,
                    int_dtype,
                );

                let output_tensor = prep.finish(
                    (
                        x_state,
                        output.indices.clone(),
                        kernel_size,
                        stride,
                        padding,
                        dilation,
                        ceil_mode,
                    ),
                    output.output,
                );

                MaxPool1dWithIndices::new(output_tensor, output.indices)
            }
            OpsKind::UnTracked(prep) => {
                let output = B::max_pool1d_with_indices(
                    x.primitive,
                    kernel_size,
                    stride,
                    padding,
                    dilation,
                    ceil_mode,
                    int_dtype,
                );
                let output_tensor = prep.finish(output.output);

                MaxPool1dWithIndices::new(output_tensor, output.indices)
            }
        }
    }

    fn max_pool1d_with_indices_backward(
        x: AutodiffTensor<B>,
        kernel_size: usize,
        stride: usize,
        padding: usize,
        dilation: usize,
        ceil_mode: bool,
        output_grad: AutodiffTensor<B>,
        indices: IntTensor<B>,
    ) -> MaxPool1dBackward<Self> {
        let output = B::max_pool1d_with_indices_backward(
            x.primitive,
            kernel_size,
            stride,
            padding,
            dilation,
            ceil_mode,
            output_grad.primitive,
            indices,
        );
        MaxPool1dBackward::new(AutodiffTensor::new(output.x_grad))
    }

    fn max_pool2d(
        x: AutodiffTensor<B>,
        kernel_size: [usize; 2],
        stride: [usize; 2],
        padding: [usize; 2],
        dilation: [usize; 2],
        ceil_mode: bool,
    ) -> AutodiffTensor<B> {
        match MaxPool2D
            .prepare::<C>([x.node.clone()])
            .compute_bound()
            .stateful()
        {
            OpsKind::Tracked(mut prep) => {
                let x_state = prep.checkpoint(&x);
                let settings = get_device_settings::<B>(&B::float_device(&x.primitive));
                let output = B::max_pool2d_with_indices(
                    x.primitive,
                    kernel_size,
                    stride,
                    padding,
                    dilation,
                    ceil_mode,
                    settings.int_dtype,
                );
                prep.finish(
                    (
                        x_state,
                        output.indices,
                        kernel_size,
                        stride,
                        padding,
                        dilation,
                        ceil_mode,
                    ),
                    output.output,
                )
            }
            OpsKind::UnTracked(prep) => prep.finish(B::max_pool2d(
                x.primitive,
                kernel_size,
                stride,
                padding,
                dilation,
                ceil_mode,
            )),
        }
    }

    fn max_pool2d_with_indices(
        x: AutodiffTensor<B>,
        kernel_size: [usize; 2],
        stride: [usize; 2],
        padding: [usize; 2],
        dilation: [usize; 2],
        ceil_mode: bool,
        int_dtype: IntDType,
    ) -> MaxPool2dWithIndices<Self> {
        match MaxPool2D
            .prepare::<C>([x.node.clone()])
            .compute_bound()
            .stateful()
        {
            OpsKind::Tracked(mut prep) => {
                let x_state = prep.checkpoint(&x);

                let output = B::max_pool2d_with_indices(
                    x.primitive,
                    kernel_size,
                    stride,
                    padding,
                    dilation,
                    ceil_mode,
                    int_dtype,
                );

                let output_tensor = prep.finish(
                    (
                        x_state,
                        output.indices.clone(),
                        kernel_size,
                        stride,
                        padding,
                        dilation,
                        ceil_mode,
                    ),
                    output.output,
                );

                MaxPool2dWithIndices::new(output_tensor, output.indices)
            }
            OpsKind::UnTracked(prep) => {
                let output = B::max_pool2d_with_indices(
                    x.primitive,
                    kernel_size,
                    stride,
                    padding,
                    dilation,
                    ceil_mode,
                    int_dtype,
                );
                let output_tensor = prep.finish(output.output);

                MaxPool2dWithIndices::new(output_tensor, output.indices)
            }
        }
    }

    fn max_pool2d_with_indices_backward(
        _x: AutodiffTensor<B>,
        _kernel_size: [usize; 2],
        _stride: [usize; 2],
        _padding: [usize; 2],
        _dilation: [usize; 2],
        _ceil_mode: bool,
        _output_grad: AutodiffTensor<B>,
        _indices: IntTensor<B>,
    ) -> MaxPool2dBackward<Self> {
        panic!("Can't differentiate max pool2d with indices backward.");
    }
    fn adaptive_avg_pool1d(x: AutodiffTensor<B>, output_size: usize) -> AutodiffTensor<B> {
        #[derive(Debug)]
        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,
            ) {
                let [node_parent] = ops.parents;
                let grad = grads.consume::<B>(&ops.node);
                let state = checkpointer.retrieve_node_output(ops.state);

                if let Some(node) = node_parent {
                    let grad = B::adaptive_avg_pool1d_backward(state, grad);
                    grads.register::<B>(node.id, grad);
                }
            }
        }

        match AdaptiveAvgPool1D
            .prepare::<C>([x.node.clone()])
            .compute_bound()
            .stateful()
        {
            OpsKind::Tracked(mut prep) => {
                let x_state = prep.checkpoint(&x);
                prep.finish(x_state, B::adaptive_avg_pool1d(x.primitive, output_size))
            }
            OpsKind::UnTracked(prep) => {
                prep.finish(B::adaptive_avg_pool1d(x.primitive, output_size))
            }
        }
    }

    fn adaptive_avg_pool2d(x: AutodiffTensor<B>, output_size: [usize; 2]) -> AutodiffTensor<B> {
        #[derive(Debug)]
        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,
            ) {
                let [node_parent] = ops.parents;
                let grad = grads.consume::<B>(&ops.node);
                let state = checkpointer.retrieve_node_output(ops.state);

                if let Some(node) = node_parent {
                    let grad = B::adaptive_avg_pool2d_backward(state, grad);
                    grads.register::<B>(node.id, grad);
                }
            }
        }

        match AdaptiveAvgPool2D
            .prepare::<C>([x.node.clone()])
            .compute_bound()
            .stateful()
        {
            OpsKind::Tracked(mut prep) => {
                let x_state = prep.checkpoint(&x);
                prep.finish(x_state, B::adaptive_avg_pool2d(x.primitive, output_size))
            }
            OpsKind::UnTracked(prep) => {
                prep.finish(B::adaptive_avg_pool2d(x.primitive, output_size))
            }
        }
    }

    fn adaptive_avg_pool2d_backward(
        _x: AutodiffTensor<B>,
        _grad: AutodiffTensor<B>,
    ) -> AutodiffTensor<B> {
        panic!("Can't differentiate adaptive avg pool2d backward.");
    }

    fn interpolate(
        x: AutodiffTensor<B>,
        output_size: [usize; 2],
        options: InterpolateOptions,
    ) -> AutodiffTensor<B> {
        #[derive(Debug)]
        struct Interpolate;
        impl<B: Backend> Backward<B, 1> for Interpolate {
            type State = (NodeId, [usize; 2], InterpolateOptions);

            fn backward(
                self,
                ops: Ops<Self::State, 1>,
                grads: &mut Gradients,
                checkpointer: &mut Checkpointer,
            ) {
                let [node_parent] = ops.parents;
                let grad = grads.consume::<B>(&ops.node);

                let (x_state, output_size, options) = ops.state;
                let state = checkpointer.retrieve_node_output(x_state);

                if let Some(node) = node_parent {
                    let grad = B::interpolate_backward(state, grad, output_size, options);
                    grads.register::<B>(node.id, grad);
                }
            }
        }

        match Interpolate
            .prepare::<C>([x.node.clone()])
            .compute_bound()
            .stateful()
        {
            OpsKind::Tracked(mut prep) => {
                let x_state = prep.checkpoint(&x);
                let output = B::interpolate(x.primitive.clone(), output_size, options.clone());
                prep.finish((x_state, output_size, options), output)
            }
            OpsKind::UnTracked(prep) => {
                prep.finish(B::interpolate(x.primitive, output_size, options))
            }
        }
    }

    fn interpolate_backward(
        _x: FloatTensor<Autodiff<B, C>>,
        _grad: FloatTensor<Autodiff<B, C>>,
        _output_size: [usize; 2],
        _options: InterpolateOptions,
    ) -> AutodiffTensor<B> {
        panic!("Can't differentiate interpolate backward.");
    }

    fn attention(
        query: FloatTensor<Autodiff<B, C>>,
        key: FloatTensor<Autodiff<B, C>>,
        value: FloatTensor<Autodiff<B, C>>,
        mask: Option<burn_backend::tensor::BoolTensor<Autodiff<B, C>>>,
        attn_bias: Option<FloatTensor<Autodiff<B, C>>>,
        options: AttentionModuleOptions,
    ) -> FloatTensor<Autodiff<B, C>> {
        attention_fallback::<Self>(query, key, value, mask, attn_bias, options)
    }

    fn ctc_loss(
        log_probs: FloatTensor<Autodiff<B, C>>,
        targets: IntTensor<Autodiff<B, C>>,
        input_lengths: IntTensor<Autodiff<B, C>>,
        target_lengths: IntTensor<Autodiff<B, C>>,
        blank: usize,
    ) -> FloatTensor<Autodiff<B, C>> {
        // Backends without a native ctc_loss_backward fall back to the default
        // implementation, which is built from differentiable tensor ops so the
        // autodiff layer derives the gradient automatically.
        if !B::has_ctc_loss_backward() {
            return burn_backend::ops::ctc::ctc_loss_default::<Self>(
                log_probs,
                targets,
                input_lengths,
                target_lengths,
                blank,
            );
        }

        #[derive(Debug)]
        struct CtcLoss;

        impl<B: Backend> Backward<B, 1> for CtcLoss {
            type State = (NodeId, IntTensor<B>, IntTensor<B>, IntTensor<B>, usize);

            fn backward(
                self,
                ops: Ops<Self::State, 1>,
                grads: &mut Gradients,
                checkpointer: &mut Checkpointer,
            ) {
                let [node_parent] = ops.parents;
                let grad_loss = grads.consume::<B>(&ops.node);

                let (log_probs_state, targets, input_lengths, target_lengths, blank) = ops.state;
                let log_probs: B::FloatTensorPrimitive =
                    checkpointer.retrieve_node_output(log_probs_state);

                if let Some(node) = node_parent {
                    let grad = B::ctc_loss_backward(
                        log_probs,
                        targets,
                        input_lengths,
                        target_lengths,
                        grad_loss,
                        blank,
                    );
                    grads.register::<B>(node.id, grad);
                }
            }
        }

        match CtcLoss
            .prepare::<C>([log_probs.node.clone()])
            .compute_bound()
            .stateful()
        {
            OpsKind::Tracked(mut prep) => {
                let log_probs_state = prep.checkpoint(&log_probs);
                let output = B::ctc_loss(
                    log_probs.primitive.clone(),
                    targets.clone(),
                    input_lengths.clone(),
                    target_lengths.clone(),
                    blank,
                );
                prep.finish(
                    (
                        log_probs_state,
                        targets,
                        input_lengths,
                        target_lengths,
                        blank,
                    ),
                    output,
                )
            }
            OpsKind::UnTracked(prep) => prep.finish(B::ctc_loss(
                log_probs.primitive,
                targets,
                input_lengths,
                target_lengths,
                blank,
            )),
        }
    }

    fn rfft(
        signal: FloatTensor<Autodiff<B, C>>,
        dim: usize,
        n: Option<usize>,
    ) -> (FloatTensor<Autodiff<B, C>>, FloatTensor<Autodiff<B, C>>) {
        #[derive(Debug)]
        struct Rfft;

        impl<B: Backend> Backward<B, 1> for Rfft {
            type State = (usize, Option<usize>, usize, usize, Vec<Slice>, Vec<Slice>);

            fn backward(
                self,
                ops: Ops<Self::State, 1>,
                grads: &mut Gradients,
                _checkpointer: &mut Checkpointer,
            ) {
                unary::<B, _>(ops.parents, ops.node, grads, |grad| {
                    let (dim, n, input_len, n_fft, slices_re, slices_im) = ops.state;

                    let grad_re = B::float_slice(grad.clone(), &slices_re);
                    let grad_im = B::float_slice(grad.clone(), &slices_im);

                    let grad_re = mul_interior::<B>(grad_re, dim, n_fft, 0.5);
                    let grad_im = mul_interior::<B>(grad_im, dim, n_fft, 0.5);

                    let grad = B::irfft(grad_re, grad_im, dim, n);
                    let grad = B::float_mul_scalar(grad, (n_fft as f64).into());

                    pad_to_length::<B>(grad, dim, input_len)
                });
            }
        }

        let input_len = signal.shape()[dim];
        let n_fft = n.unwrap_or(input_len);
        let (re, im) = B::rfft(signal.primitive, dim, n);

        // In order to perform only a single irfft in the backward pass, we have to temporarily
        // bundle `re` and `im` into a single tensor. The following slice vecs are used to split
        // them back up in both the forward and the backward pass.
        let slices_re = re
            .shape()
            .iter()
            .map(|&len| Slice::from(0..len))
            .collect::<Vec<Slice>>();
        let slices_im = {
            let mut slices = slices_re.clone();
            let len = slices[0].end.unwrap();
            slices[0].start = len;
            slices[0].end = Some(2 * len);
            slices
        };
        let spectrum = B::float_cat(vec![re, im], 0);
        let state = (
            dim,
            n,
            input_len,
            n_fft,
            slices_re.clone(),
            slices_im.clone(),
        );

        let spectrum = match Rfft
            .prepare::<C>([signal.node.clone()])
            .compute_bound()
            .stateful()
        {
            OpsKind::Tracked(prep) => prep.finish(state, spectrum),
            OpsKind::UnTracked(prep) => prep.finish(spectrum),
        };

        let re = Self::float_slice(spectrum.clone(), &slices_re);
        let im = Self::float_slice(spectrum.clone(), &slices_im);

        (re, im)
    }

    fn irfft(
        spectrum_re: FloatTensor<Autodiff<B, C>>,
        spectrum_im: FloatTensor<Autodiff<B, C>>,
        dim: usize,
        n: Option<usize>,
    ) -> FloatTensor<Autodiff<B, C>> {
        #[derive(Debug)]
        struct Irfft;

        impl<B: Backend> Backward<B, 2> for Irfft {
            type State = (usize, Option<usize>, usize, usize);

            fn backward(
                self,
                ops: Ops<Self::State, 2>,
                grads: &mut Gradients,
                _checkpointer: &mut Checkpointer,
            ) {
                let (dim, n, input_len, n_fft) = ops.state;
                let [node_re, node_im] = ops.parents;
                let grad = grads.consume::<B>(&ops.node);

                let grad = B::float_div_scalar(grad, (n_fft as f64).into());

                let (grad_re, grad_im) = B::rfft(grad, dim, n);

                let grad_re = mul_interior::<B>(grad_re, dim, n_fft, 2.0);
                let grad_im = mul_interior::<B>(grad_im, dim, n_fft, 2.0);

                let grad_re = pad_to_length::<B>(grad_re, dim, input_len);
                let grad_im = pad_to_length::<B>(grad_im, dim, input_len);

                if let Some(node) = node_re {
                    grads.register::<B>(node.id, grad_re);
                }
                if let Some(node) = node_im {
                    grads.register::<B>(node.id, grad_im);
                }
            }
        }

        let input_len = spectrum_re.shape()[dim];
        let signal = B::irfft(spectrum_re.primitive, spectrum_im.primitive, dim, n);
        let n_fft = n.unwrap_or(signal.shape()[dim]);
        let state = (dim, n, input_len, n_fft);

        match Irfft
            .prepare::<C>([spectrum_re.node.clone(), spectrum_im.node.clone()])
            .compute_bound()
            .stateful()
        {
            OpsKind::Tracked(prep) => prep.finish(state, signal),
            OpsKind::UnTracked(prep) => prep.finish(signal),
        }
    }
}

// adapted from: burn_cubecl::kernel::fft::base::pad_to_length
fn pad_to_length<B: Backend>(tensor: FloatTensor<B>, dim: usize, target: usize) -> FloatTensor<B> {
    let shape = tensor.shape();
    let current = shape[dim];
    if current == target {
        return tensor;
    }
    if current > target {
        let slices: Vec<_> = shape
            .iter()
            .enumerate()
            .map(|(i, &s)| Slice::from(if i == dim { 0..target } else { 0..s }))
            .collect();
        return B::float_slice(tensor, &slices);
    }
    let mut padded_shape = shape.clone();
    padded_shape[dim] = target;
    let padded = B::float_zeros(
        padded_shape,
        &B::float_device(&tensor),
        tensor.dtype().into(),
    );
    let slices: Vec<Slice> = shape.iter().map(|&s| Slice::from(0..s)).collect();
    B::float_slice_assign(padded, &slices, tensor)
}

fn mul_interior<B: Backend>(
    bins: FloatTensor<B>,
    dim: usize,
    n_fft: usize,
    factor: f64,
) -> FloatTensor<B> {
    // identify the interior bins (all bins except DC and Nyquist)
    let slices_interior: Vec<Slice> = {
        let mut ranges = bins.shape().into_ranges();

        // skip the DC bin
        ranges[dim].start += 1;

        // if `n_fft` is even, we have a Nyquist bin to skip
        if n_fft.is_multiple_of(2) {
            ranges[dim].end -= 1;
        }

        ranges.into_iter().map(Slice::from).collect()
    };

    // multiply only the interior bins by `factor`
    let interior = B::float_slice(bins.clone(), &slices_interior);
    let interior = B::float_mul_scalar(interior, factor.into());

    B::float_slice_assign(bins, &slices_interior, interior)
}

#[derive(Debug)]
struct MaxPool1D;

impl<B: Backend> Backward<B, 1> for MaxPool1D {
    type State = (NodeId, IntTensor<B>, usize, usize, usize, usize, bool);

    fn backward(
        self,
        ops: Ops<Self::State, 1>,
        grads: &mut Gradients,
        checkpointer: &mut Checkpointer,
    ) {
        let [node_parent] = ops.parents;
        let grad = grads.consume::<B>(&ops.node);
        let (x_state, indices, kernel_size, stride, padding, dilation, ceil_mode) = ops.state;
        let x = checkpointer.retrieve_node_output(x_state);

        if let Some(node) = node_parent {
            let grad = B::max_pool1d_with_indices_backward(
                x,
                kernel_size,
                stride,
                padding,
                dilation,
                ceil_mode,
                grad,
                indices,
            );

            grads.register::<B>(node.id, grad.x_grad);
        }
    }
}

#[derive(Debug)]
struct MaxPool2D;

impl<B: Backend> Backward<B, 1> for MaxPool2D {
    type State = (
        NodeId,
        IntTensor<B>,
        [usize; 2],
        [usize; 2],
        [usize; 2],
        [usize; 2],
        bool,
    );

    fn backward(
        self,
        ops: Ops<Self::State, 1>,
        grads: &mut Gradients,
        checkpointer: &mut Checkpointer,
    ) {
        let [node_parent] = ops.parents;
        let grad = grads.consume::<B>(&ops.node);
        let (x_state, indices, kernel_size, stride, padding, dilation, ceil_mode) = ops.state;
        let x = checkpointer.retrieve_node_output(x_state);

        if let Some(node) = node_parent {
            let grad = B::max_pool2d_with_indices_backward(
                x,
                kernel_size,
                stride,
                padding,
                dilation,
                ceil_mode,
                grad,
                indices,
            );

            grads.register::<B>(node.id, grad.x_grad);
        }
    }
}
