Trait ModuleOps 

pub trait ModuleOps<B>
where
    B: Backend,
{
    // Required methods
    fn conv2d(
        x: <B as BackendTypes>::FloatTensorPrimitive,
        weight: <B as BackendTypes>::FloatTensorPrimitive,
        bias: Option<<B as BackendTypes>::FloatTensorPrimitive>,
        options: ConvOptions<2>,
    ) -> <B as BackendTypes>::FloatTensorPrimitive;
    fn deform_conv2d(
        x: <B as BackendTypes>::FloatTensorPrimitive,
        offset: <B as BackendTypes>::FloatTensorPrimitive,
        weight: <B as BackendTypes>::FloatTensorPrimitive,
        mask: Option<<B as BackendTypes>::FloatTensorPrimitive>,
        bias: Option<<B as BackendTypes>::FloatTensorPrimitive>,
        options: DeformConvOptions<2>,
    ) -> <B as BackendTypes>::FloatTensorPrimitive;
    fn deform_conv2d_backward(
        x: <B as BackendTypes>::FloatTensorPrimitive,
        offset: <B as BackendTypes>::FloatTensorPrimitive,
        weight: <B as BackendTypes>::FloatTensorPrimitive,
        mask: Option<<B as BackendTypes>::FloatTensorPrimitive>,
        bias: Option<<B as BackendTypes>::FloatTensorPrimitive>,
        output_grad: <B as BackendTypes>::FloatTensorPrimitive,
        options: DeformConvOptions<2>,
    ) -> DeformConv2dBackward<B>;
    fn conv3d(
        x: <B as BackendTypes>::FloatTensorPrimitive,
        weight: <B as BackendTypes>::FloatTensorPrimitive,
        bias: Option<<B as BackendTypes>::FloatTensorPrimitive>,
        options: ConvOptions<3>,
    ) -> <B as BackendTypes>::FloatTensorPrimitive;
    fn conv_transpose2d(
        x: <B as BackendTypes>::FloatTensorPrimitive,
        weight: <B as BackendTypes>::FloatTensorPrimitive,
        bias: Option<<B as BackendTypes>::FloatTensorPrimitive>,
        options: ConvTransposeOptions<2>,
    ) -> <B as BackendTypes>::FloatTensorPrimitive;
    fn conv_transpose3d(
        x: <B as BackendTypes>::FloatTensorPrimitive,
        weight: <B as BackendTypes>::FloatTensorPrimitive,
        bias: Option<<B as BackendTypes>::FloatTensorPrimitive>,
        options: ConvTransposeOptions<3>,
    ) -> <B as BackendTypes>::FloatTensorPrimitive;
    fn avg_pool2d(
        x: <B as BackendTypes>::FloatTensorPrimitive,
        kernel_size: [usize; 2],
        stride: [usize; 2],
        padding: [usize; 2],
        count_include_pad: bool,
        ceil_mode: bool,
    ) -> <B as BackendTypes>::FloatTensorPrimitive;
    fn avg_pool2d_backward(
        x: <B as BackendTypes>::FloatTensorPrimitive,
        grad: <B as BackendTypes>::FloatTensorPrimitive,
        kernel_size: [usize; 2],
        stride: [usize; 2],
        padding: [usize; 2],
        count_include_pad: bool,
        ceil_mode: bool,
    ) -> <B as BackendTypes>::FloatTensorPrimitive;
    fn adaptive_avg_pool2d(
        x: <B as BackendTypes>::FloatTensorPrimitive,
        output_size: [usize; 2],
    ) -> <B as BackendTypes>::FloatTensorPrimitive;
    fn adaptive_avg_pool2d_backward(
        x: <B as BackendTypes>::FloatTensorPrimitive,
        grad: <B as BackendTypes>::FloatTensorPrimitive,
    ) -> <B as BackendTypes>::FloatTensorPrimitive;
    fn max_pool2d(
        x: <B as BackendTypes>::FloatTensorPrimitive,
        kernel_size: [usize; 2],
        stride: [usize; 2],
        padding: [usize; 2],
        dilation: [usize; 2],
        ceil_mode: bool,
    ) -> <B as BackendTypes>::FloatTensorPrimitive;
    fn max_pool2d_with_indices(
        x: <B as BackendTypes>::FloatTensorPrimitive,
        kernel_size: [usize; 2],
        stride: [usize; 2],
        padding: [usize; 2],
        dilation: [usize; 2],
        ceil_mode: bool,
    ) -> MaxPool2dWithIndices<B>;
    fn max_pool2d_with_indices_backward(
        x: <B as BackendTypes>::FloatTensorPrimitive,
        kernel_size: [usize; 2],
        stride: [usize; 2],
        padding: [usize; 2],
        dilation: [usize; 2],
        ceil_mode: bool,
        output_grad: <B as BackendTypes>::FloatTensorPrimitive,
        indices: <B as BackendTypes>::IntTensorPrimitive,
    ) -> MaxPool2dBackward<B>;
    fn interpolate(
        x: <B as BackendTypes>::FloatTensorPrimitive,
        output_size: [usize; 2],
        options: InterpolateOptions,
    ) -> <B as BackendTypes>::FloatTensorPrimitive;
    fn interpolate_backward(
        x: <B as BackendTypes>::FloatTensorPrimitive,
        grad: <B as BackendTypes>::FloatTensorPrimitive,
        output_size: [usize; 2],
        options: InterpolateOptions,
    ) -> <B as BackendTypes>::FloatTensorPrimitive;
    fn attention(
        query: <B as BackendTypes>::FloatTensorPrimitive,
        key: <B as BackendTypes>::FloatTensorPrimitive,
        value: <B as BackendTypes>::FloatTensorPrimitive,
        mask: Option<<B as BackendTypes>::BoolTensorPrimitive>,
        attn_bias: Option<<B as BackendTypes>::FloatTensorPrimitive>,
        options: AttentionModuleOptions,
    ) -> <B as BackendTypes>::FloatTensorPrimitive;
    fn rfft(
        signal: <B as BackendTypes>::FloatTensorPrimitive,
        dim: usize,
        n: Option<usize>,
    ) -> (<B as BackendTypes>::FloatTensorPrimitive, <B as BackendTypes>::FloatTensorPrimitive);
    fn irfft(
        spectrum_re: <B as BackendTypes>::FloatTensorPrimitive,
        spectrum_im: <B as BackendTypes>::FloatTensorPrimitive,
        dim: usize,
        n: Option<usize>,
    ) -> <B as BackendTypes>::FloatTensorPrimitive;

    // Provided methods
    fn exponential_relu_native(
        tensor: <B as BackendTypes>::FloatTensorPrimitive,
        alpha: f64,
        continuous: bool,
    ) -> <B as BackendTypes>::FloatTensorPrimitive { ... }
    fn exponential_relu_native_backward(
        tensor: <B as BackendTypes>::FloatTensorPrimitive,
        grad: <B as BackendTypes>::FloatTensorPrimitive,
        alpha: f64,
        continuous: bool,
    ) -> <B as BackendTypes>::FloatTensorPrimitive { ... }
    fn leaky_relu_native(
        tensor: <B as BackendTypes>::FloatTensorPrimitive,
        negative_slope: f64,
    ) -> <B as BackendTypes>::FloatTensorPrimitive { ... }
    fn leaky_relu_native_backward(
        tensor: <B as BackendTypes>::FloatTensorPrimitive,
        grad: <B as BackendTypes>::FloatTensorPrimitive,
        negative_slope: f64,
    ) -> <B as BackendTypes>::FloatTensorPrimitive { ... }
    fn prelu_native(
        tensor: <B as BackendTypes>::FloatTensorPrimitive,
        alpha: <B as BackendTypes>::FloatTensorPrimitive,
    ) -> <B as BackendTypes>::FloatTensorPrimitive { ... }
    fn prelu_native_backward_select(
        tensor: <B as BackendTypes>::FloatTensorPrimitive,
        alpha: <B as BackendTypes>::FloatTensorPrimitive,
        grad: <B as BackendTypes>::FloatTensorPrimitive,
        mask: [bool; 2],
    ) -> [Option<<B as BackendTypes>::FloatTensorPrimitive>; 2] { ... }
    fn group_norm(
        tensor: <B as BackendTypes>::FloatTensorPrimitive,
        gamma: Option<<B as BackendTypes>::FloatTensorPrimitive>,
        beta: Option<<B as BackendTypes>::FloatTensorPrimitive>,
        groups: usize,
        epsilon: f64,
    ) -> <B as BackendTypes>::FloatTensorPrimitive { ... }
    fn group_norm_with_stats(
        tensor: <B as BackendTypes>::FloatTensorPrimitive,
        gamma: Option<<B as BackendTypes>::FloatTensorPrimitive>,
        beta: Option<<B as BackendTypes>::FloatTensorPrimitive>,
        groups: usize,
        epsilon: f64,
    ) -> LayerNormOutput<B> { ... }
    fn group_norm_backward_select(
        tensor: <B as BackendTypes>::FloatTensorPrimitive,
        gamma: Option<<B as BackendTypes>::FloatTensorPrimitive>,
        grad: <B as BackendTypes>::FloatTensorPrimitive,
        mean: <B as BackendTypes>::FloatTensorPrimitive,
        rstd: <B as BackendTypes>::FloatTensorPrimitive,
        groups: usize,
        mask: [bool; 3],
    ) -> [Option<<B as BackendTypes>::FloatTensorPrimitive>; 3] { ... }
    fn gelu_native(
        tensor: <B as BackendTypes>::FloatTensorPrimitive,
        approximate: bool,
    ) -> <B as BackendTypes>::FloatTensorPrimitive { ... }
    fn gelu_native_backward(
        input: <B as BackendTypes>::FloatTensorPrimitive,
        grad: <B as BackendTypes>::FloatTensorPrimitive,
        approximate: bool,
    ) -> <B as BackendTypes>::FloatTensorPrimitive { ... }
    fn silu_native(
        tensor: <B as BackendTypes>::FloatTensorPrimitive,
    ) -> <B as BackendTypes>::FloatTensorPrimitive { ... }
    fn silu_native_backward(
        input: <B as BackendTypes>::FloatTensorPrimitive,
        grad: <B as BackendTypes>::FloatTensorPrimitive,
    ) -> <B as BackendTypes>::FloatTensorPrimitive { ... }
    fn softmax_native(
        tensor: <B as BackendTypes>::FloatTensorPrimitive,
        dim: usize,
        logarithmic: bool,
    ) -> <B as BackendTypes>::FloatTensorPrimitive { ... }
    fn softmax_with_stats(
        tensor: <B as BackendTypes>::FloatTensorPrimitive,
        dim: usize,
        logarithmic: bool,
    ) -> SoftmaxOutput<B> { ... }
    fn softmax_native_backward(
        working: <B as BackendTypes>::FloatTensorPrimitive,
        grad: <B as BackendTypes>::FloatTensorPrimitive,
        dim: usize,
        logarithmic: bool,
    ) -> <B as BackendTypes>::FloatTensorPrimitive { ... }
    fn embedding(
        weights: <B as BackendTypes>::FloatTensorPrimitive,
        indices: <B as BackendTypes>::IntTensorPrimitive,
    ) -> <B as BackendTypes>::FloatTensorPrimitive { ... }
    fn embedding_backward(
        weights: <B as BackendTypes>::FloatTensorPrimitive,
        output_grad: <B as BackendTypes>::FloatTensorPrimitive,
        indices: <B as BackendTypes>::IntTensorPrimitive,
    ) -> <B as BackendTypes>::FloatTensorPrimitive { ... }
    fn linear(
        x: <B as BackendTypes>::FloatTensorPrimitive,
        weight: <B as BackendTypes>::FloatTensorPrimitive,
        bias: Option<<B as BackendTypes>::FloatTensorPrimitive>,
    ) -> <B as BackendTypes>::FloatTensorPrimitive { ... }
    fn linear_x_backward(
        weight: <B as BackendTypes>::FloatTensorPrimitive,
        output_grad: <B as BackendTypes>::FloatTensorPrimitive,
    ) -> <B as BackendTypes>::FloatTensorPrimitive { ... }
    fn linear_weight_backward(
        x: <B as BackendTypes>::FloatTensorPrimitive,
        output_grad: <B as BackendTypes>::FloatTensorPrimitive,
    ) -> <B as BackendTypes>::FloatTensorPrimitive { ... }
    fn linear_bias_backward(
        output_grad: <B as BackendTypes>::FloatTensorPrimitive,
    ) -> <B as BackendTypes>::FloatTensorPrimitive { ... }
    fn conv1d(
        x: <B as BackendTypes>::FloatTensorPrimitive,
        weight: <B as BackendTypes>::FloatTensorPrimitive,
        bias: Option<<B as BackendTypes>::FloatTensorPrimitive>,
        options: ConvOptions<1>,
    ) -> <B as BackendTypes>::FloatTensorPrimitive { ... }
    fn conv1d_x_backward(
        x: <B as BackendTypes>::FloatTensorPrimitive,
        weight: <B as BackendTypes>::FloatTensorPrimitive,
        output_grad: <B as BackendTypes>::FloatTensorPrimitive,
        options: ConvOptions<1>,
    ) -> <B as BackendTypes>::FloatTensorPrimitive { ... }
    fn conv1d_weight_backward(
        x: <B as BackendTypes>::FloatTensorPrimitive,
        weight: <B as BackendTypes>::FloatTensorPrimitive,
        output_grad: <B as BackendTypes>::FloatTensorPrimitive,
        options: ConvOptions<1>,
    ) -> <B as BackendTypes>::FloatTensorPrimitive { ... }
    fn conv1d_bias_backward(
        x: <B as BackendTypes>::FloatTensorPrimitive,
        bias: <B as BackendTypes>::FloatTensorPrimitive,
        output_grad: <B as BackendTypes>::FloatTensorPrimitive,
    ) -> <B as BackendTypes>::FloatTensorPrimitive { ... }
    fn conv2d_x_backward(
        x: <B as BackendTypes>::FloatTensorPrimitive,
        weight: <B as BackendTypes>::FloatTensorPrimitive,
        output_grad: <B as BackendTypes>::FloatTensorPrimitive,
        options: ConvOptions<2>,
    ) -> <B as BackendTypes>::FloatTensorPrimitive { ... }
    fn conv2d_weight_backward(
        x: <B as BackendTypes>::FloatTensorPrimitive,
        weight: <B as BackendTypes>::FloatTensorPrimitive,
        output_grad: <B as BackendTypes>::FloatTensorPrimitive,
        options: ConvOptions<2>,
    ) -> <B as BackendTypes>::FloatTensorPrimitive { ... }
    fn conv2d_bias_backward(
        x: <B as BackendTypes>::FloatTensorPrimitive,
        bias: <B as BackendTypes>::FloatTensorPrimitive,
        output_grad: <B as BackendTypes>::FloatTensorPrimitive,
    ) -> <B as BackendTypes>::FloatTensorPrimitive { ... }
    fn conv3d_x_backward(
        x: <B as BackendTypes>::FloatTensorPrimitive,
        weight: <B as BackendTypes>::FloatTensorPrimitive,
        output_grad: <B as BackendTypes>::FloatTensorPrimitive,
        options: ConvOptions<3>,
    ) -> <B as BackendTypes>::FloatTensorPrimitive { ... }
    fn conv3d_weight_backward(
        x: <B as BackendTypes>::FloatTensorPrimitive,
        weight: <B as BackendTypes>::FloatTensorPrimitive,
        output_grad: <B as BackendTypes>::FloatTensorPrimitive,
        options: ConvOptions<3>,
    ) -> <B as BackendTypes>::FloatTensorPrimitive { ... }
    fn conv3d_bias_backward(
        x: <B as BackendTypes>::FloatTensorPrimitive,
        bias: <B as BackendTypes>::FloatTensorPrimitive,
        output_grad: <B as BackendTypes>::FloatTensorPrimitive,
    ) -> <B as BackendTypes>::FloatTensorPrimitive { ... }
    fn conv_transpose1d(
        x: <B as BackendTypes>::FloatTensorPrimitive,
        weight: <B as BackendTypes>::FloatTensorPrimitive,
        bias: Option<<B as BackendTypes>::FloatTensorPrimitive>,
        options: ConvTransposeOptions<1>,
    ) -> <B as BackendTypes>::FloatTensorPrimitive { ... }
    fn conv_transpose1d_x_backward(
        weight: <B as BackendTypes>::FloatTensorPrimitive,
        output_grad: <B as BackendTypes>::FloatTensorPrimitive,
        options: ConvTransposeOptions<1>,
    ) -> <B as BackendTypes>::FloatTensorPrimitive { ... }
    fn conv_transpose1d_weight_backward(
        x: <B as BackendTypes>::FloatTensorPrimitive,
        weight: <B as BackendTypes>::FloatTensorPrimitive,
        output_grad: <B as BackendTypes>::FloatTensorPrimitive,
        options: ConvTransposeOptions<1>,
    ) -> <B as BackendTypes>::FloatTensorPrimitive { ... }
    fn conv_transpose1d_bias_backward(
        x: <B as BackendTypes>::FloatTensorPrimitive,
        bias: <B as BackendTypes>::FloatTensorPrimitive,
        output_grad: <B as BackendTypes>::FloatTensorPrimitive,
    ) -> <B as BackendTypes>::FloatTensorPrimitive { ... }
    fn conv_transpose2d_x_backward(
        weight: <B as BackendTypes>::FloatTensorPrimitive,
        output_grad: <B as BackendTypes>::FloatTensorPrimitive,
        options: ConvTransposeOptions<2>,
    ) -> <B as BackendTypes>::FloatTensorPrimitive { ... }
    fn conv_transpose2d_weight_backward(
        x: <B as BackendTypes>::FloatTensorPrimitive,
        weight: <B as BackendTypes>::FloatTensorPrimitive,
        output_grad: <B as BackendTypes>::FloatTensorPrimitive,
        options: ConvTransposeOptions<2>,
    ) -> <B as BackendTypes>::FloatTensorPrimitive { ... }
    fn conv_transpose2d_bias_backward(
        x: <B as BackendTypes>::FloatTensorPrimitive,
        bias: <B as BackendTypes>::FloatTensorPrimitive,
        output_grad: <B as BackendTypes>::FloatTensorPrimitive,
    ) -> <B as BackendTypes>::FloatTensorPrimitive { ... }
    fn conv_transpose3d_x_backward(
        weight: <B as BackendTypes>::FloatTensorPrimitive,
        output_grad: <B as BackendTypes>::FloatTensorPrimitive,
        options: ConvTransposeOptions<3>,
    ) -> <B as BackendTypes>::FloatTensorPrimitive { ... }
    fn conv_transpose3d_weight_backward(
        x: <B as BackendTypes>::FloatTensorPrimitive,
        weight: <B as BackendTypes>::FloatTensorPrimitive,
        output_grad: <B as BackendTypes>::FloatTensorPrimitive,
        options: ConvTransposeOptions<3>,
    ) -> <B as BackendTypes>::FloatTensorPrimitive { ... }
    fn conv_transpose3d_bias_backward(
        x: <B as BackendTypes>::FloatTensorPrimitive,
        bias: <B as BackendTypes>::FloatTensorPrimitive,
        output_grad: <B as BackendTypes>::FloatTensorPrimitive,
    ) -> <B as BackendTypes>::FloatTensorPrimitive { ... }
    fn unfold4d(
        x: <B as BackendTypes>::FloatTensorPrimitive,
        kernel_size: [usize; 2],
        options: UnfoldOptions,
    ) -> <B as BackendTypes>::FloatTensorPrimitive { ... }
    fn avg_pool1d(
        x: <B as BackendTypes>::FloatTensorPrimitive,
        kernel_size: usize,
        stride: usize,
        padding: usize,
        count_include_pad: bool,
        ceil_mode: bool,
    ) -> <B as BackendTypes>::FloatTensorPrimitive { ... }
    fn avg_pool1d_backward(
        x: <B as BackendTypes>::FloatTensorPrimitive,
        grad: <B as BackendTypes>::FloatTensorPrimitive,
        kernel_size: usize,
        stride: usize,
        padding: usize,
        count_include_pad: bool,
        ceil_mode: bool,
    ) -> <B as BackendTypes>::FloatTensorPrimitive { ... }
    fn avg_pool3d(
        x: <B as BackendTypes>::FloatTensorPrimitive,
        kernel_size: [usize; 3],
        stride: [usize; 3],
        padding: [usize; 3],
        count_include_pad: bool,
        ceil_mode: bool,
    ) -> <B as BackendTypes>::FloatTensorPrimitive { ... }
    fn avg_pool3d_native_output_size(
        _input: [usize; 3],
        _kernel: [usize; 3],
        _stride: [usize; 3],
        _padding: [usize; 3],
        _ceil: bool,
    ) -> Option<[usize; 3]> { ... }
    fn avg_pool3d_backward(
        x: <B as BackendTypes>::FloatTensorPrimitive,
        grad: <B as BackendTypes>::FloatTensorPrimitive,
        kernel_size: [usize; 3],
        stride: [usize; 3],
        padding: [usize; 3],
        count_include_pad: bool,
        ceil_mode: bool,
    ) -> <B as BackendTypes>::FloatTensorPrimitive { ... }
    fn adaptive_avg_pool3d(
        x: <B as BackendTypes>::FloatTensorPrimitive,
        output_size: [usize; 3],
    ) -> <B as BackendTypes>::FloatTensorPrimitive { ... }
    fn adaptive_avg_pool3d_backward(
        x: <B as BackendTypes>::FloatTensorPrimitive,
        grad: <B as BackendTypes>::FloatTensorPrimitive,
    ) -> <B as BackendTypes>::FloatTensorPrimitive { ... }
    fn max_pool3d(
        x: <B as BackendTypes>::FloatTensorPrimitive,
        kernel: [usize; 3],
        stride: [usize; 3],
        padding: [usize; 3],
        dilation: [usize; 3],
        ceil: bool,
    ) -> <B as BackendTypes>::FloatTensorPrimitive { ... }
    fn max_pool3d_with_indices(
        x: <B as BackendTypes>::FloatTensorPrimitive,
        kernel: [usize; 3],
        stride: [usize; 3],
        padding: [usize; 3],
        dilation: [usize; 3],
        ceil: bool,
    ) -> MaxPool3dWithIndices<B> { ... }
    fn max_pool3d_with_indices_backward(
        x: <B as BackendTypes>::FloatTensorPrimitive,
        grad: <B as BackendTypes>::FloatTensorPrimitive,
        indices: <B as BackendTypes>::IntTensorPrimitive,
        _kernel: [usize; 3],
        _stride: [usize; 3],
        _padding: [usize; 3],
        _dilation: [usize; 3],
        _ceil: bool,
    ) -> MaxPool3dBackward<B> { ... }
    fn adaptive_avg_pool1d(
        x: <B as BackendTypes>::FloatTensorPrimitive,
        output_size: usize,
    ) -> <B as BackendTypes>::FloatTensorPrimitive { ... }
    fn adaptive_avg_pool1d_backward(
        x: <B as BackendTypes>::FloatTensorPrimitive,
        grad: <B as BackendTypes>::FloatTensorPrimitive,
    ) -> <B as BackendTypes>::FloatTensorPrimitive { ... }
    fn max_pool1d(
        x: <B as BackendTypes>::FloatTensorPrimitive,
        kernel_size: usize,
        stride: usize,
        padding: usize,
        dilation: usize,
        ceil_mode: bool,
    ) -> <B as BackendTypes>::FloatTensorPrimitive { ... }
    fn max_pool1d_with_indices(
        x: <B as BackendTypes>::FloatTensorPrimitive,
        kernel_size: usize,
        stride: usize,
        padding: usize,
        dilation: usize,
        ceil_mode: bool,
    ) -> MaxPool1dWithIndices<B> { ... }
    fn max_pool1d_with_indices_backward(
        x: <B as BackendTypes>::FloatTensorPrimitive,
        kernel_size: usize,
        stride: usize,
        padding: usize,
        dilation: usize,
        ceil_mode: bool,
        output_grad: <B as BackendTypes>::FloatTensorPrimitive,
        indices: <B as BackendTypes>::IntTensorPrimitive,
    ) -> MaxPool1dBackward<B> { ... }
    fn interpolate1d(
        x: <B as BackendTypes>::FloatTensorPrimitive,
        size: usize,
        options: InterpolateOptions,
    ) -> <B as BackendTypes>::FloatTensorPrimitive { ... }
    fn interpolate1d_backward(
        x: <B as BackendTypes>::FloatTensorPrimitive,
        grad: <B as BackendTypes>::FloatTensorPrimitive,
        size: usize,
        options: InterpolateOptions,
    ) -> <B as BackendTypes>::FloatTensorPrimitive { ... }
    fn interpolate3d(
        x: <B as BackendTypes>::FloatTensorPrimitive,
        size: [usize; 3],
        options: InterpolateOptions,
    ) -> <B as BackendTypes>::FloatTensorPrimitive { ... }
    fn interpolate3d_backward(
        x: <B as BackendTypes>::FloatTensorPrimitive,
        grad: <B as BackendTypes>::FloatTensorPrimitive,
        size: [usize; 3],
        options: InterpolateOptions,
    ) -> <B as BackendTypes>::FloatTensorPrimitive { ... }
    fn layer_norm(
        tensor: <B as BackendTypes>::FloatTensorPrimitive,
        gamma: <B as BackendTypes>::FloatTensorPrimitive,
        beta: Option<<B as BackendTypes>::FloatTensorPrimitive>,
        epsilon: f64,
    ) -> <B as BackendTypes>::FloatTensorPrimitive { ... }
    fn has_layer_norm_backward() -> bool { ... }
    fn has_rms_norm_backward() -> bool { ... }
    fn rms_norm(
        tensor: <B as BackendTypes>::FloatTensorPrimitive,
        gamma: <B as BackendTypes>::FloatTensorPrimitive,
        epsilon: f64,
    ) -> <B as BackendTypes>::FloatTensorPrimitive { ... }
    fn rms_norm_with_stats(
        tensor: <B as BackendTypes>::FloatTensorPrimitive,
        gamma: <B as BackendTypes>::FloatTensorPrimitive,
        epsilon: f64,
    ) -> RmsNormOutput<B> { ... }
    fn rms_norm_backward(
        tensor: <B as BackendTypes>::FloatTensorPrimitive,
        gamma: <B as BackendTypes>::FloatTensorPrimitive,
        grad: <B as BackendTypes>::FloatTensorPrimitive,
        rstd: <B as BackendTypes>::FloatTensorPrimitive,
    ) -> RmsNormBackward<B> { ... }
    fn rms_norm_backward_select(
        tensor: <B as BackendTypes>::FloatTensorPrimitive,
        gamma: <B as BackendTypes>::FloatTensorPrimitive,
        grad: <B as BackendTypes>::FloatTensorPrimitive,
        rstd: <B as BackendTypes>::FloatTensorPrimitive,
        mask: [bool; 2],
    ) -> [Option<<B as BackendTypes>::FloatTensorPrimitive>; 2] { ... }
    fn layer_norm_with_stats(
        tensor: <B as BackendTypes>::FloatTensorPrimitive,
        gamma: <B as BackendTypes>::FloatTensorPrimitive,
        beta: Option<<B as BackendTypes>::FloatTensorPrimitive>,
        epsilon: f64,
    ) -> LayerNormOutput<B> { ... }
    fn layer_norm_backward(
        tensor: <B as BackendTypes>::FloatTensorPrimitive,
        gamma: <B as BackendTypes>::FloatTensorPrimitive,
        grad: <B as BackendTypes>::FloatTensorPrimitive,
        mean: <B as BackendTypes>::FloatTensorPrimitive,
        rstd: <B as BackendTypes>::FloatTensorPrimitive,
    ) -> LayerNormBackward<B> { ... }
    fn layer_norm_backward_select(
        tensor: <B as BackendTypes>::FloatTensorPrimitive,
        gamma: <B as BackendTypes>::FloatTensorPrimitive,
        grad: <B as BackendTypes>::FloatTensorPrimitive,
        mean: <B as BackendTypes>::FloatTensorPrimitive,
        rstd: <B as BackendTypes>::FloatTensorPrimitive,
        mask: [bool; 3],
    ) -> [Option<<B as BackendTypes>::FloatTensorPrimitive>; 3] { ... }
    fn layer_norm_default(
        tensor: <B as BackendTypes>::FloatTensorPrimitive,
        gamma: <B as BackendTypes>::FloatTensorPrimitive,
        beta: Option<<B as BackendTypes>::FloatTensorPrimitive>,
        epsilon: f64,
    ) -> <B as BackendTypes>::FloatTensorPrimitive { ... }
    fn ctc_loss(
        log_probs: <B as BackendTypes>::FloatTensorPrimitive,
        targets: <B as BackendTypes>::IntTensorPrimitive,
        input_lengths: <B as BackendTypes>::IntTensorPrimitive,
        target_lengths: <B as BackendTypes>::IntTensorPrimitive,
        blank: usize,
    ) -> <B as BackendTypes>::FloatTensorPrimitive { ... }
    fn has_ctc_loss_backward() -> bool { ... }
    fn ctc_loss_backward(
        _log_probs: <B as BackendTypes>::FloatTensorPrimitive,
        _targets: <B as BackendTypes>::IntTensorPrimitive,
        _input_lengths: <B as BackendTypes>::IntTensorPrimitive,
        _target_lengths: <B as BackendTypes>::IntTensorPrimitive,
        _grad_loss: <B as BackendTypes>::FloatTensorPrimitive,
        _blank: usize,
    ) -> <B as BackendTypes>::FloatTensorPrimitive { ... }
}

Module operations trait.

Required Methods

fn conv2d( x: <B as BackendTypes>::FloatTensorPrimitive, weight: <B as BackendTypes>::FloatTensorPrimitive, bias: Option<<B as BackendTypes>::FloatTensorPrimitive>, options: ConvOptions<2>, ) -> <B as BackendTypes>::FloatTensorPrimitive

Two dimensional convolution.

Shapes

x: [batch_size, channels_in, height, width], weight: [channels_out, channels_in, kernel_size_1, kernel_size_2], bias: [channels_out],

fn deform_conv2d( x: <B as BackendTypes>::FloatTensorPrimitive, offset: <B as BackendTypes>::FloatTensorPrimitive, weight: <B as BackendTypes>::FloatTensorPrimitive, mask: Option<<B as BackendTypes>::FloatTensorPrimitive>, bias: Option<<B as BackendTypes>::FloatTensorPrimitive>, options: DeformConvOptions<2>, ) -> <B as BackendTypes>::FloatTensorPrimitive

Two dimensional deformable convolution.

Shapes

x: [batch_size, channels_in, height, width], weight: [channels_out, channels_in, kernel_size_1, kernel_size_2], bias: [channels_out],

fn deform_conv2d_backward( x: <B as BackendTypes>::FloatTensorPrimitive, offset: <B as BackendTypes>::FloatTensorPrimitive, weight: <B as BackendTypes>::FloatTensorPrimitive, mask: Option<<B as BackendTypes>::FloatTensorPrimitive>, bias: Option<<B as BackendTypes>::FloatTensorPrimitive>, output_grad: <B as BackendTypes>::FloatTensorPrimitive, options: DeformConvOptions<2>, ) -> DeformConv2dBackward<B>

Backward pass for the deform_conv2d operation.

fn conv3d( x: <B as BackendTypes>::FloatTensorPrimitive, weight: <B as BackendTypes>::FloatTensorPrimitive, bias: Option<<B as BackendTypes>::FloatTensorPrimitive>, options: ConvOptions<3>, ) -> <B as BackendTypes>::FloatTensorPrimitive

Three dimensional convolution.

Shapes

x: [batch_size, channels_in, depth, height, width], weight: [channels_out, channels_in, kernel_size_1, kernel_size_2, kernel_size_3], bias: [channels_out],

fn conv_transpose2d( x: <B as BackendTypes>::FloatTensorPrimitive, weight: <B as BackendTypes>::FloatTensorPrimitive, bias: Option<<B as BackendTypes>::FloatTensorPrimitive>, options: ConvTransposeOptions<2>, ) -> <B as BackendTypes>::FloatTensorPrimitive

Two dimensional transposed convolution.

Shapes

x: [batch_size, channels_in, height, width], weight: [channels_in, channels_out, kernel_size_1, kernel_size_2], bias: [channels_out],

fn conv_transpose3d( x: <B as BackendTypes>::FloatTensorPrimitive, weight: <B as BackendTypes>::FloatTensorPrimitive, bias: Option<<B as BackendTypes>::FloatTensorPrimitive>, options: ConvTransposeOptions<3>, ) -> <B as BackendTypes>::FloatTensorPrimitive

Three dimensional transposed convolution.

Shapes

x: [batch_size, channels_in, height, width], weight: [channels_in, channels_out, kernel_size_1, kernel_size_2, kernel_size_3], bias: [channels_out],

fn avg_pool2d( x: <B as BackendTypes>::FloatTensorPrimitive, kernel_size: [usize; 2], stride: [usize; 2], padding: [usize; 2], count_include_pad: bool, ceil_mode: bool, ) -> <B as BackendTypes>::FloatTensorPrimitive

Two dimensional avg pooling.

Shapes

x: [batch_size, channels, height, width],

fn avg_pool2d_backward( x: <B as BackendTypes>::FloatTensorPrimitive, grad: <B as BackendTypes>::FloatTensorPrimitive, kernel_size: [usize; 2], stride: [usize; 2], padding: [usize; 2], count_include_pad: bool, ceil_mode: bool, ) -> <B as BackendTypes>::FloatTensorPrimitive

Backward pass for the avg pooling 2d operation.

fn adaptive_avg_pool2d( x: <B as BackendTypes>::FloatTensorPrimitive, output_size: [usize; 2], ) -> <B as BackendTypes>::FloatTensorPrimitive

Two dimensional adaptive avg pooling.

Shapes

x: [batch_size, channels, height, width],

fn adaptive_avg_pool2d_backward( x: <B as BackendTypes>::FloatTensorPrimitive, grad: <B as BackendTypes>::FloatTensorPrimitive, ) -> <B as BackendTypes>::FloatTensorPrimitive

Backward pass for the adaptive avg pooling 2d operation.

fn max_pool2d( x: <B as BackendTypes>::FloatTensorPrimitive, kernel_size: [usize; 2], stride: [usize; 2], padding: [usize; 2], dilation: [usize; 2], ceil_mode: bool, ) -> <B as BackendTypes>::FloatTensorPrimitive

Two dimensional max pooling.

Shapes

x: [batch_size, channels, height, width],

fn max_pool2d_with_indices( x: <B as BackendTypes>::FloatTensorPrimitive, kernel_size: [usize; 2], stride: [usize; 2], padding: [usize; 2], dilation: [usize; 2], ceil_mode: bool, ) -> MaxPool2dWithIndices<B>

Two dimensional max pooling with indices.

Shapes

x: [batch_size, channels, height, width],

fn max_pool2d_with_indices_backward( x: <B as BackendTypes>::FloatTensorPrimitive, kernel_size: [usize; 2], stride: [usize; 2], padding: [usize; 2], dilation: [usize; 2], ceil_mode: bool, output_grad: <B as BackendTypes>::FloatTensorPrimitive, indices: <B as BackendTypes>::IntTensorPrimitive, ) -> MaxPool2dBackward<B>

Backward pass for the max pooling 2d operation.

fn interpolate( x: <B as BackendTypes>::FloatTensorPrimitive, output_size: [usize; 2], options: InterpolateOptions, ) -> <B as BackendTypes>::FloatTensorPrimitive

Down/up samples the input.

Shapes

x: [batch_size, channels, height, width],

fn interpolate_backward( x: <B as BackendTypes>::FloatTensorPrimitive, grad: <B as BackendTypes>::FloatTensorPrimitive, output_size: [usize; 2], options: InterpolateOptions, ) -> <B as BackendTypes>::FloatTensorPrimitive

Backward pass for the interpolate operation.

fn attention( query: <B as BackendTypes>::FloatTensorPrimitive, key: <B as BackendTypes>::FloatTensorPrimitive, value: <B as BackendTypes>::FloatTensorPrimitive, mask: Option<<B as BackendTypes>::BoolTensorPrimitive>, attn_bias: Option<<B as BackendTypes>::FloatTensorPrimitive>, options: AttentionModuleOptions, ) -> <B as BackendTypes>::FloatTensorPrimitive

Computes scaled dot-product attention: softmax(QKᵗ * scale) · V, where scale defaults to 1/sqrt(head_dim). Optionally applies masking, additive bias, causal masking, and softcap to the attention scores.

Arguments

• query: Query tensor of shape [batch_size, num_heads, seq_len_q, head_dim]
• key: Key tensor of shape [batch_size, num_heads, seq_len_k, head_dim]
• value: Value tensor of shape [batch_size, num_heads, seq_len_k, val_dim]
• mask: Optional boolean mask of shape [batch_size, num_heads, seq_len_q, seq_len_k], where true indicates positions to mask (i.e. set to -inf before softmax).
• attn_bias: Optional float tensor of shape [batch_size, num_heads, seq_len_q, seq_len_k] added to the attention scores before softmax (e.g. ALiBi, relative position biases).
• options: Additional attention options (custom scale, softcap, causal masking).

Returns

A tensor of shape [batch_size, num_heads, seq_len_q, val_dim] representing the attended context per head.

Note

This implementation does not support dropout and is intended for inference or use cases where dropout is not needed.

fn rfft( signal: <B as BackendTypes>::FloatTensorPrimitive, dim: usize, n: Option<usize>, ) -> (<B as BackendTypes>::FloatTensorPrimitive, <B as BackendTypes>::FloatTensorPrimitive)

Real-valued FFT with optional size parameter.

When n is None, the signal must be a power of two along dim, and the output has signal_len / 2 + 1 frequency bins.

When n is Some(size), size must also be a power of two. The signal is truncated or zero-padded to size and the output has size / 2 + 1 frequency bins. Non-power- of-two sizes are currently rejected at the public API boundary; true arbitrary-n DFT support (Bluestein’s algorithm) is tracked as a follow-up.

Returns two tensors: the real part and the imaginary part.

fn irfft( spectrum_re: <B as BackendTypes>::FloatTensorPrimitive, spectrum_im: <B as BackendTypes>::FloatTensorPrimitive, dim: usize, n: Option<usize>, ) -> <B as BackendTypes>::FloatTensorPrimitive

Inverse real-valued FFT with optional output size.

When n is None, the reconstructed signal length 2 * (spectrum_size - 1) must be a power of two.

When n is Some(size), size must also be a power of two. Output has exactly size samples.

Provided Methods

fn exponential_relu_native( tensor: <B as BackendTypes>::FloatTensorPrimitive, alpha: f64, continuous: bool, ) -> <B as BackendTypes>::FloatTensorPrimitive

Working-storage ELU, or CELU when continuous is true, using the original alpha.

fn exponential_relu_native_backward( tensor: <B as BackendTypes>::FloatTensorPrimitive, grad: <B as BackendTypes>::FloatTensorPrimitive, alpha: f64, continuous: bool, ) -> <B as BackendTypes>::FloatTensorPrimitive

Original selected ELU/CELU mode’s independent first-order input derivative.

fn leaky_relu_native( tensor: <B as BackendTypes>::FloatTensorPrimitive, negative_slope: f64, ) -> <B as BackendTypes>::FloatTensorPrimitive

Working-storage LeakyReLU using the actual scalar negative slope.

fn leaky_relu_native_backward( tensor: <B as BackendTypes>::FloatTensorPrimitive, grad: <B as BackendTypes>::FloatTensorPrimitive, negative_slope: f64, ) -> <B as BackendTypes>::FloatTensorPrimitive

Independent first-order VJP, using the original primal to select the branch.

fn prelu_native( tensor: <B as BackendTypes>::FloatTensorPrimitive, alpha: <B as BackendTypes>::FloatTensorPrimitive, ) -> <B as BackendTypes>::FloatTensorPrimitive

Working-storage PReLU with the original shared or per-channel slope vector.

fn prelu_native_backward_select( tensor: <B as BackendTypes>::FloatTensorPrimitive, alpha: <B as BackendTypes>::FloatTensorPrimitive, grad: <B as BackendTypes>::FloatTensorPrimitive, mask: [bool; 2], ) -> [Option<<B as BackendTypes>::FloatTensorPrimitive>; 2]

Independently requested original-input and slope derivatives.

fn group_norm( tensor: <B as BackendTypes>::FloatTensorPrimitive, gamma: Option<<B as BackendTypes>::FloatTensorPrimitive>, beta: Option<<B as BackendTypes>::FloatTensorPrimitive>, groups: usize, epsilon: f64, ) -> <B as BackendTypes>::FloatTensorPrimitive

Native GroupNorm on actual channel groups with independently optional affine leaves.

fn group_norm_with_stats( tensor: <B as BackendTypes>::FloatTensorPrimitive, gamma: Option<<B as BackendTypes>::FloatTensorPrimitive>, beta: Option<<B as BackendTypes>::FloatTensorPrimitive>, groups: usize, epsilon: f64, ) -> LayerNormOutput<B>

GroupNorm output and FP32/FP64 saved [batch, groups] mean/reciprocal deviation.

fn group_norm_backward_select( tensor: <B as BackendTypes>::FloatTensorPrimitive, gamma: Option<<B as BackendTypes>::FloatTensorPrimitive>, grad: <B as BackendTypes>::FloatTensorPrimitive, mean: <B as BackendTypes>::FloatTensorPrimitive, rstd: <B as BackendTypes>::FloatTensorPrimitive, groups: usize, mask: [bool; 3], ) -> [Option<<B as BackendTypes>::FloatTensorPrimitive>; 3]

Only requested original input, actual weight and bias derivatives, in that order.

fn gelu_native( tensor: <B as BackendTypes>::FloatTensorPrimitive, approximate: bool, ) -> <B as BackendTypes>::FloatTensorPrimitive

Explicit native GELU preserving the caller’s original erf/tanh approximation choice.

fn gelu_native_backward( input: <B as BackendTypes>::FloatTensorPrimitive, grad: <B as BackendTypes>::FloatTensorPrimitive, approximate: bool, ) -> <B as BackendTypes>::FloatTensorPrimitive

Independent first-order GELU VJP in the same selected erf/tanh mode.

fn silu_native( tensor: <B as BackendTypes>::FloatTensorPrimitive, ) -> <B as BackendTypes>::FloatTensorPrimitive

Explicit native SiLU training path, using FP32 for half activation arithmetic.

fn silu_native_backward( input: <B as BackendTypes>::FloatTensorPrimitive, grad: <B as BackendTypes>::FloatTensorPrimitive, ) -> <B as BackendTypes>::FloatTensorPrimitive

Independent first-order SiLU VJP, retaining the original input storage.

fn softmax_native( tensor: <B as BackendTypes>::FloatTensorPrimitive, dim: usize, logarithmic: bool, ) -> <B as BackendTypes>::FloatTensorPrimitive

Explicit saved-working-output softmax path; existing activation defaults are unchanged.

fn softmax_with_stats( tensor: <B as BackendTypes>::FloatTensorPrimitive, dim: usize, logarithmic: bool, ) -> SoftmaxOutput<B>

Softmax/log-softmax retaining its working output, with FP32 low-precision arithmetic.

fn softmax_native_backward( working: <B as BackendTypes>::FloatTensorPrimitive, grad: <B as BackendTypes>::FloatTensorPrimitive, dim: usize, logarithmic: bool, ) -> <B as BackendTypes>::FloatTensorPrimitive

VJP from the saved working output, retaining FP64 if either argument uses FP64.

fn embedding( weights: <B as BackendTypes>::FloatTensorPrimitive, indices: <B as BackendTypes>::IntTensorPrimitive, ) -> <B as BackendTypes>::FloatTensorPrimitive

Embedding operation.

Arguments

• weights - The embedding weights.
• indices - The indices tensor.

Returns

The output tensor.

fn embedding_backward( weights: <B as BackendTypes>::FloatTensorPrimitive, output_grad: <B as BackendTypes>::FloatTensorPrimitive, indices: <B as BackendTypes>::IntTensorPrimitive, ) -> <B as BackendTypes>::FloatTensorPrimitive

Embedding backward operation.

Arguments

• weights - The embedding weights.
• output_grad - The output gradient.
• indices - The indices tensor.

Returns

The gradient.

fn linear( x: <B as BackendTypes>::FloatTensorPrimitive, weight: <B as BackendTypes>::FloatTensorPrimitive, bias: Option<<B as BackendTypes>::FloatTensorPrimitive>, ) -> <B as BackendTypes>::FloatTensorPrimitive

Linear transformation.

Shapes

x: [..., d_input], weight: [d_input, d_output], bias: [d_output],

fn linear_x_backward( weight: <B as BackendTypes>::FloatTensorPrimitive, output_grad: <B as BackendTypes>::FloatTensorPrimitive, ) -> <B as BackendTypes>::FloatTensorPrimitive

Backward pass for linear, returning the gradient for x.

fn linear_weight_backward( x: <B as BackendTypes>::FloatTensorPrimitive, output_grad: <B as BackendTypes>::FloatTensorPrimitive, ) -> <B as BackendTypes>::FloatTensorPrimitive

Backward pass for linear, returning the gradient for weight.

fn linear_bias_backward( output_grad: <B as BackendTypes>::FloatTensorPrimitive, ) -> <B as BackendTypes>::FloatTensorPrimitive

Backward pass for linear, returning the gradient for bias.

fn conv1d( x: <B as BackendTypes>::FloatTensorPrimitive, weight: <B as BackendTypes>::FloatTensorPrimitive, bias: Option<<B as BackendTypes>::FloatTensorPrimitive>, options: ConvOptions<1>, ) -> <B as BackendTypes>::FloatTensorPrimitive

One dimensional convolution.

Shapes

x: [batch_size, channels_in, length], weight: [channels_out, channels_in, kernel_size], bias: [channels_out],

fn conv1d_x_backward( x: <B as BackendTypes>::FloatTensorPrimitive, weight: <B as BackendTypes>::FloatTensorPrimitive, output_grad: <B as BackendTypes>::FloatTensorPrimitive, options: ConvOptions<1>, ) -> <B as BackendTypes>::FloatTensorPrimitive

Backward pass for the conv1d operation, returning the gradient for x.

fn conv1d_weight_backward( x: <B as BackendTypes>::FloatTensorPrimitive, weight: <B as BackendTypes>::FloatTensorPrimitive, output_grad: <B as BackendTypes>::FloatTensorPrimitive, options: ConvOptions<1>, ) -> <B as BackendTypes>::FloatTensorPrimitive

Backward pass for the conv1d operation, returning the gradient for weight.

fn conv1d_bias_backward( x: <B as BackendTypes>::FloatTensorPrimitive, bias: <B as BackendTypes>::FloatTensorPrimitive, output_grad: <B as BackendTypes>::FloatTensorPrimitive, ) -> <B as BackendTypes>::FloatTensorPrimitive

Backward pass for the conv1d operation, returning the gradient for bias.

fn conv2d_x_backward( x: <B as BackendTypes>::FloatTensorPrimitive, weight: <B as BackendTypes>::FloatTensorPrimitive, output_grad: <B as BackendTypes>::FloatTensorPrimitive, options: ConvOptions<2>, ) -> <B as BackendTypes>::FloatTensorPrimitive

Backward pass for the conv2d operation, returning the gradient for x.

fn conv2d_weight_backward( x: <B as BackendTypes>::FloatTensorPrimitive, weight: <B as BackendTypes>::FloatTensorPrimitive, output_grad: <B as BackendTypes>::FloatTensorPrimitive, options: ConvOptions<2>, ) -> <B as BackendTypes>::FloatTensorPrimitive

Backward pass for the conv2d operation, returning the gradient for weight.

fn conv2d_bias_backward( x: <B as BackendTypes>::FloatTensorPrimitive, bias: <B as BackendTypes>::FloatTensorPrimitive, output_grad: <B as BackendTypes>::FloatTensorPrimitive, ) -> <B as BackendTypes>::FloatTensorPrimitive

Backward pass for the conv2d operation, returning the gradient for bias.

fn conv3d_x_backward( x: <B as BackendTypes>::FloatTensorPrimitive, weight: <B as BackendTypes>::FloatTensorPrimitive, output_grad: <B as BackendTypes>::FloatTensorPrimitive, options: ConvOptions<3>, ) -> <B as BackendTypes>::FloatTensorPrimitive

Backward pass for the conv3d operation, returning the gradient for x.

fn conv3d_weight_backward( x: <B as BackendTypes>::FloatTensorPrimitive, weight: <B as BackendTypes>::FloatTensorPrimitive, output_grad: <B as BackendTypes>::FloatTensorPrimitive, options: ConvOptions<3>, ) -> <B as BackendTypes>::FloatTensorPrimitive

Backward pass for the conv3d operation, returning the gradient for weight.

fn conv3d_bias_backward( x: <B as BackendTypes>::FloatTensorPrimitive, bias: <B as BackendTypes>::FloatTensorPrimitive, output_grad: <B as BackendTypes>::FloatTensorPrimitive, ) -> <B as BackendTypes>::FloatTensorPrimitive

Backward pass for the conv3d operation, returning the gradient for bias.

fn conv_transpose1d( x: <B as BackendTypes>::FloatTensorPrimitive, weight: <B as BackendTypes>::FloatTensorPrimitive, bias: Option<<B as BackendTypes>::FloatTensorPrimitive>, options: ConvTransposeOptions<1>, ) -> <B as BackendTypes>::FloatTensorPrimitive

One dimensional transposed convolution.

Shapes

x: [batch_size, channels_in, length], weight: [channels_in, channels_out, length], bias: [channels_out],

fn conv_transpose1d_x_backward( weight: <B as BackendTypes>::FloatTensorPrimitive, output_grad: <B as BackendTypes>::FloatTensorPrimitive, options: ConvTransposeOptions<1>, ) -> <B as BackendTypes>::FloatTensorPrimitive

Backward pass for the conv transpose 1d operation, returning the gradient for x.

fn conv_transpose1d_weight_backward( x: <B as BackendTypes>::FloatTensorPrimitive, weight: <B as BackendTypes>::FloatTensorPrimitive, output_grad: <B as BackendTypes>::FloatTensorPrimitive, options: ConvTransposeOptions<1>, ) -> <B as BackendTypes>::FloatTensorPrimitive

Backward pass for the conv transpose 1d operation, returning the gradient for weight.

fn conv_transpose1d_bias_backward( x: <B as BackendTypes>::FloatTensorPrimitive, bias: <B as BackendTypes>::FloatTensorPrimitive, output_grad: <B as BackendTypes>::FloatTensorPrimitive, ) -> <B as BackendTypes>::FloatTensorPrimitive

Backward pass for the conv transpose 1d operation, returning the gradient for bias.

fn conv_transpose2d_x_backward( weight: <B as BackendTypes>::FloatTensorPrimitive, output_grad: <B as BackendTypes>::FloatTensorPrimitive, options: ConvTransposeOptions<2>, ) -> <B as BackendTypes>::FloatTensorPrimitive

Backward pass for the conv transpose 2d operation, returning the gradient for x.

fn conv_transpose2d_weight_backward( x: <B as BackendTypes>::FloatTensorPrimitive, weight: <B as BackendTypes>::FloatTensorPrimitive, output_grad: <B as BackendTypes>::FloatTensorPrimitive, options: ConvTransposeOptions<2>, ) -> <B as BackendTypes>::FloatTensorPrimitive

Backward pass for the conv transpose 2d operation, returning the gradient for weight.

fn conv_transpose2d_bias_backward( x: <B as BackendTypes>::FloatTensorPrimitive, bias: <B as BackendTypes>::FloatTensorPrimitive, output_grad: <B as BackendTypes>::FloatTensorPrimitive, ) -> <B as BackendTypes>::FloatTensorPrimitive

Backward pass for the conv transpose 2d operation, returning the gradient for bias.

fn conv_transpose3d_x_backward( weight: <B as BackendTypes>::FloatTensorPrimitive, output_grad: <B as BackendTypes>::FloatTensorPrimitive, options: ConvTransposeOptions<3>, ) -> <B as BackendTypes>::FloatTensorPrimitive

Backward pass for the conv transpose 3d operation, returning the gradient for x.

fn conv_transpose3d_weight_backward( x: <B as BackendTypes>::FloatTensorPrimitive, weight: <B as BackendTypes>::FloatTensorPrimitive, output_grad: <B as BackendTypes>::FloatTensorPrimitive, options: ConvTransposeOptions<3>, ) -> <B as BackendTypes>::FloatTensorPrimitive

Backward pass for the conv transpose 3d operation, returning the gradient for weight.

fn conv_transpose3d_bias_backward( x: <B as BackendTypes>::FloatTensorPrimitive, bias: <B as BackendTypes>::FloatTensorPrimitive, output_grad: <B as BackendTypes>::FloatTensorPrimitive, ) -> <B as BackendTypes>::FloatTensorPrimitive

Backward pass for the conv transpose 3d operation, returning the gradient for bias.

fn unfold4d( x: <B as BackendTypes>::FloatTensorPrimitive, kernel_size: [usize; 2], options: UnfoldOptions, ) -> <B as BackendTypes>::FloatTensorPrimitive

Four-dimensional unfolding.

Shapes

• x: [batch_size, channels_in, height, width],
• returns: [batch_size, channels_in * kernel_size_1 * kernel_size_2, number of blocks],

fn avg_pool1d( x: <B as BackendTypes>::FloatTensorPrimitive, kernel_size: usize, stride: usize, padding: usize, count_include_pad: bool, ceil_mode: bool, ) -> <B as BackendTypes>::FloatTensorPrimitive

One dimensional avg pooling.

Shapes

x: [batch_size, channels, length],

fn avg_pool1d_backward( x: <B as BackendTypes>::FloatTensorPrimitive, grad: <B as BackendTypes>::FloatTensorPrimitive, kernel_size: usize, stride: usize, padding: usize, count_include_pad: bool, ceil_mode: bool, ) -> <B as BackendTypes>::FloatTensorPrimitive

Backward pass for the avg pooling 1d operation.

fn avg_pool3d( x: <B as BackendTypes>::FloatTensorPrimitive, kernel_size: [usize; 3], stride: [usize; 3], padding: [usize; 3], count_include_pad: bool, ceil_mode: bool, ) -> <B as BackendTypes>::FloatTensorPrimitive

Three dimensional average pooling of [batch, channels, depth, height, width].

fn avg_pool3d_native_output_size( _input: [usize; 3], _kernel: [usize; 3], _stride: [usize; 3], _padding: [usize; 3], _ceil: bool, ) -> Option<[usize; 3]>

Native volume output geometry, or None to retain the backend’s composed operations.

fn avg_pool3d_backward( x: <B as BackendTypes>::FloatTensorPrimitive, grad: <B as BackendTypes>::FloatTensorPrimitive, kernel_size: [usize; 3], stride: [usize; 3], padding: [usize; 3], count_include_pad: bool, ceil_mode: bool, ) -> <B as BackendTypes>::FloatTensorPrimitive

Input gradients for three dimensional average pooling.

fn adaptive_avg_pool3d( x: <B as BackendTypes>::FloatTensorPrimitive, output_size: [usize; 3], ) -> <B as BackendTypes>::FloatTensorPrimitive

Three dimensional adaptive average pooling of [batch, channels, depth, height, width].

fn adaptive_avg_pool3d_backward( x: <B as BackendTypes>::FloatTensorPrimitive, grad: <B as BackendTypes>::FloatTensorPrimitive, ) -> <B as BackendTypes>::FloatTensorPrimitive

Input gradients for native three dimensional adaptive average pooling.

fn max_pool3d( x: <B as BackendTypes>::FloatTensorPrimitive, kernel: [usize; 3], stride: [usize; 3], padding: [usize; 3], dilation: [usize; 3], ceil: bool, ) -> <B as BackendTypes>::FloatTensorPrimitive

Three dimensional maximum pooling with native padding, stride and dilation.

fn max_pool3d_with_indices( x: <B as BackendTypes>::FloatTensorPrimitive, kernel: [usize; 3], stride: [usize; 3], padding: [usize; 3], dilation: [usize; 3], ceil: bool, ) -> MaxPool3dWithIndices<B>

Volume maxima and original input positions.

fn max_pool3d_with_indices_backward( x: <B as BackendTypes>::FloatTensorPrimitive, grad: <B as BackendTypes>::FloatTensorPrimitive, indices: <B as BackendTypes>::IntTensorPrimitive, _kernel: [usize; 3], _stride: [usize; 3], _padding: [usize; 3], _dilation: [usize; 3], _ceil: bool, ) -> MaxPool3dBackward<B>

Accumulate volume gradients through the saved maximum positions.

fn adaptive_avg_pool1d( x: <B as BackendTypes>::FloatTensorPrimitive, output_size: usize, ) -> <B as BackendTypes>::FloatTensorPrimitive

One dimensional adaptive avg pooling.

Shapes

x: [batch_size, channels, length],

fn adaptive_avg_pool1d_backward( x: <B as BackendTypes>::FloatTensorPrimitive, grad: <B as BackendTypes>::FloatTensorPrimitive, ) -> <B as BackendTypes>::FloatTensorPrimitive

Backward pass for the adaptive avg pooling 1d operation.

fn max_pool1d( x: <B as BackendTypes>::FloatTensorPrimitive, kernel_size: usize, stride: usize, padding: usize, dilation: usize, ceil_mode: bool, ) -> <B as BackendTypes>::FloatTensorPrimitive

One dimensional max pooling.

Shapes

x: [batch_size, channels, length],

fn max_pool1d_with_indices( x: <B as BackendTypes>::FloatTensorPrimitive, kernel_size: usize, stride: usize, padding: usize, dilation: usize, ceil_mode: bool, ) -> MaxPool1dWithIndices<B>

One dimensional max pooling with indices.

Shapes

x: [batch_size, channels, height, width],

fn max_pool1d_with_indices_backward( x: <B as BackendTypes>::FloatTensorPrimitive, kernel_size: usize, stride: usize, padding: usize, dilation: usize, ceil_mode: bool, output_grad: <B as BackendTypes>::FloatTensorPrimitive, indices: <B as BackendTypes>::IntTensorPrimitive, ) -> MaxPool1dBackward<B>

Backward pass for the max pooling 1d operation.

fn interpolate1d( x: <B as BackendTypes>::FloatTensorPrimitive, size: usize, options: InterpolateOptions, ) -> <B as BackendTypes>::FloatTensorPrimitive

Resize native [batch, channels, width] lines with the original filter.

fn interpolate1d_backward( x: <B as BackendTypes>::FloatTensorPrimitive, grad: <B as BackendTypes>::FloatTensorPrimitive, size: usize, options: InterpolateOptions, ) -> <B as BackendTypes>::FloatTensorPrimitive

Backward line resizing in the original input geometry and storage.

fn interpolate3d( x: <B as BackendTypes>::FloatTensorPrimitive, size: [usize; 3], options: InterpolateOptions, ) -> <B as BackendTypes>::FloatTensorPrimitive

Resize native volumes using the existing spatial-then-depth filter sequence.

fn interpolate3d_backward( x: <B as BackendTypes>::FloatTensorPrimitive, grad: <B as BackendTypes>::FloatTensorPrimitive, size: [usize; 3], options: InterpolateOptions, ) -> <B as BackendTypes>::FloatTensorPrimitive

Backward volume resizing through the original spatial and depth filters.

fn layer_norm( tensor: <B as BackendTypes>::FloatTensorPrimitive, gamma: <B as BackendTypes>::FloatTensorPrimitive, beta: Option<<B as BackendTypes>::FloatTensorPrimitive>, epsilon: f64, ) -> <B as BackendTypes>::FloatTensorPrimitive

Applies Layer Normalization over the last dimension of the input tensor.

Computes (x - mean) / sqrt(var + epsilon) * gamma + beta, where mean and (biased) var are reduced over the last axis.

Arguments

• tensor - Input tensor of shape [..., d_model].
• gamma - Scale tensor of shape [d_model].
• beta - Optional bias tensor of shape [d_model].
• epsilon - Numerical stability term added to the variance before the square root.

Returns

A tensor with the same shape as tensor.

fn has_layer_norm_backward() -> bool

Whether native forward statistics and complete first-order backward are available.

fn has_rms_norm_backward() -> bool

Whether saved RMSNorm statistics and complete first-order backward are available.

fn rms_norm( tensor: <B as BackendTypes>::FloatTensorPrimitive, gamma: <B as BackendTypes>::FloatTensorPrimitive, epsilon: f64, ) -> <B as BackendTypes>::FloatTensorPrimitive

Last-axis RMSNorm with working arithmetic and one final output-storage cast.

fn rms_norm_with_stats( tensor: <B as BackendTypes>::FloatTensorPrimitive, gamma: <B as BackendTypes>::FloatTensorPrimitive, epsilon: f64, ) -> RmsNormOutput<B>

RMSNorm forward retaining reciprocal row norms for backward.

fn rms_norm_backward( tensor: <B as BackendTypes>::FloatTensorPrimitive, gamma: <B as BackendTypes>::FloatTensorPrimitive, grad: <B as BackendTypes>::FloatTensorPrimitive, rstd: <B as BackendTypes>::FloatTensorPrimitive, ) -> RmsNormBackward<B>

RMSNorm derivatives from saved reciprocal row norms, without repeating forward.

fn rms_norm_backward_select( tensor: <B as BackendTypes>::FloatTensorPrimitive, gamma: <B as BackendTypes>::FloatTensorPrimitive, grad: <B as BackendTypes>::FloatTensorPrimitive, rstd: <B as BackendTypes>::FloatTensorPrimitive, mask: [bool; 2], ) -> [Option<<B as BackendTypes>::FloatTensorPrimitive>; 2]

RMSNorm input and weight gradients selected by the two corresponding mask entries.

fn layer_norm_with_stats( tensor: <B as BackendTypes>::FloatTensorPrimitive, gamma: <B as BackendTypes>::FloatTensorPrimitive, beta: Option<<B as BackendTypes>::FloatTensorPrimitive>, epsilon: f64, ) -> LayerNormOutput<B>

Native forward with statistics retained for backward.

fn layer_norm_backward( tensor: <B as BackendTypes>::FloatTensorPrimitive, gamma: <B as BackendTypes>::FloatTensorPrimitive, grad: <B as BackendTypes>::FloatTensorPrimitive, mean: <B as BackendTypes>::FloatTensorPrimitive, rstd: <B as BackendTypes>::FloatTensorPrimitive, ) -> LayerNormBackward<B>

Native backward using the exact statistics returned by forward.

fn layer_norm_backward_select( tensor: <B as BackendTypes>::FloatTensorPrimitive, gamma: <B as BackendTypes>::FloatTensorPrimitive, grad: <B as BackendTypes>::FloatTensorPrimitive, mean: <B as BackendTypes>::FloatTensorPrimitive, rstd: <B as BackendTypes>::FloatTensorPrimitive, mask: [bool; 3], ) -> [Option<<B as BackendTypes>::FloatTensorPrimitive>; 3]

LayerNorm input, weight and bias gradients selected in that order.

fn layer_norm_default( tensor: <B as BackendTypes>::FloatTensorPrimitive, gamma: <B as BackendTypes>::FloatTensorPrimitive, beta: Option<<B as BackendTypes>::FloatTensorPrimitive>, epsilon: f64, ) -> <B as BackendTypes>::FloatTensorPrimitive

Differentiable primitive composition for backends without native LayerNorm backward.

fn ctc_loss( log_probs: <B as BackendTypes>::FloatTensorPrimitive, targets: <B as BackendTypes>::IntTensorPrimitive, input_lengths: <B as BackendTypes>::IntTensorPrimitive, target_lengths: <B as BackendTypes>::IntTensorPrimitive, blank: usize, ) -> <B as BackendTypes>::FloatTensorPrimitive

Computes the Connectionist Temporal Classification (CTC) loss.

Sums over all valid alignments between the input and target sequences using the forward (alpha) algorithm.

Arguments

• log_probs - Log-probabilities of shape [T, N, C]
• targets - Target label indices of shape [N, S]
• input_lengths - Actual input sequence lengths per batch element [N]
• target_lengths - Actual target lengths per batch element [N]
• blank - Index of the blank label

Returns

Per-sample loss of shape [N]

fn has_ctc_loss_backward() -> bool

Returns true if this backend implements ctc_loss_backward natively.

Autodiff queries this flag to decide between two paths:

• true: use the backend’s ctc_loss and ctc_loss_backward directly.
• false: call ctc::ctc_loss_default for the forward pass; autodiff then differentiates through the decomposed tensor ops.

Backends that override ctc_loss_backward must also override this to return true.

fn ctc_loss_backward( _log_probs: <B as BackendTypes>::FloatTensorPrimitive, _targets: <B as BackendTypes>::IntTensorPrimitive, _input_lengths: <B as BackendTypes>::IntTensorPrimitive, _target_lengths: <B as BackendTypes>::IntTensorPrimitive, _grad_loss: <B as BackendTypes>::FloatTensorPrimitive, _blank: usize, ) -> <B as BackendTypes>::FloatTensorPrimitive

Backward pass for ctc_loss: gradient w.r.t. log_probs.

Only called when has_ctc_loss_backward returns true. Backends without a native implementation should leave both methods at their defaults; the gradient is computed automatically by autodiff against the decomposed ctc::ctc_loss_default forward.

Arguments

• log_probs - Log-probabilities of shape [T, N, C]
• targets - Target label indices of shape [N, S]
• input_lengths - Actual input sequence lengths per batch element [N]
• target_lengths - Actual target lengths per batch element [N]
• grad_loss - Upstream gradient w.r.t. the per-sample loss [N]
• blank - Index of the blank label

Returns

Gradient w.r.t. log_probs of shape [T, N, C]

Dyn Compatibility

This trait is not dyn compatible.

In older versions of Rust, dyn compatibility was called "object safety".

Implementors