Trait ModuleOps 

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

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

Module operations trait.

Required Methods

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

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: FloatTensor<B>, offset: FloatTensor<B>, weight: FloatTensor<B>, mask: Option<FloatTensor<B>>, bias: Option<FloatTensor<B>>, options: DeformConvOptions<2>, ) -> FloatTensor<B>

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: FloatTensor<B>, offset: FloatTensor<B>, weight: FloatTensor<B>, mask: Option<FloatTensor<B>>, bias: Option<FloatTensor<B>>, output_grad: FloatTensor<B>, options: DeformConvOptions<2>, ) -> DeformConv2dBackward<B>

Backward pass for the deform_conv2d operation.

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

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: FloatTensor<B>, weight: FloatTensor<B>, bias: Option<FloatTensor<B>>, options: ConvTransposeOptions<2>, ) -> FloatTensor<B>

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: FloatTensor<B>, weight: FloatTensor<B>, bias: Option<FloatTensor<B>>, options: ConvTransposeOptions<3>, ) -> FloatTensor<B>

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: FloatTensor<B>, kernel_size: [usize; 2], stride: [usize; 2], padding: [usize; 2], count_include_pad: bool, ceil_mode: bool, ) -> FloatTensor<B>

Two dimensional avg pooling.

Shapes

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

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

Backward pass for the avg pooling 2d operation.

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

Two dimensional adaptive avg pooling.

Shapes

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

fn adaptive_avg_pool2d_backward( x: FloatTensor<B>, grad: FloatTensor<B>, ) -> FloatTensor<B>

Backward pass for the adaptive avg pooling 2d operation.

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

Two dimensional max pooling.

Shapes

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

fn max_pool2d_with_indices( x: FloatTensor<B>, 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: FloatTensor<B>, kernel_size: [usize; 2], stride: [usize; 2], padding: [usize; 2], dilation: [usize; 2], ceil_mode: bool, output_grad: FloatTensor<B>, indices: IntTensor<B>, ) -> MaxPool2dBackward<B>

Backward pass for the max pooling 2d operation.

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

Down/up samples the input.

Shapes

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

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

Backward pass for the interpolate operation.

fn attention( query: FloatTensor<B>, key: FloatTensor<B>, value: FloatTensor<B>, mask: Option<BoolTensor<B>>, attn_bias: Option<FloatTensor<B>>, options: AttentionModuleOptions, ) -> FloatTensor<B>

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: FloatTensor<B>, dim: usize, n: Option<usize>, ) -> (FloatTensor<B>, FloatTensor<B>)

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: FloatTensor<B>, spectrum_im: FloatTensor<B>, dim: usize, n: Option<usize>, ) -> FloatTensor<B>

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: FloatTensor<B>, alpha: f64, continuous: bool, ) -> FloatTensor<B>

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

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

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

fn leaky_relu_native( tensor: FloatTensor<B>, negative_slope: f64, ) -> FloatTensor<B>

Working-storage LeakyReLU using the actual scalar negative slope.

fn leaky_relu_native_backward( tensor: FloatTensor<B>, grad: FloatTensor<B>, negative_slope: f64, ) -> FloatTensor<B>

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

fn prelu_native(tensor: FloatTensor<B>, alpha: FloatTensor<B>) -> FloatTensor<B>

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

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

Independently requested original-input and slope derivatives.

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

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

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

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

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

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

fn gelu_native(tensor: FloatTensor<B>, approximate: bool) -> FloatTensor<B>

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

fn gelu_native_backward( input: FloatTensor<B>, grad: FloatTensor<B>, approximate: bool, ) -> FloatTensor<B>

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

fn silu_native(tensor: FloatTensor<B>) -> FloatTensor<B>

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

fn silu_native_backward( input: FloatTensor<B>, grad: FloatTensor<B>, ) -> FloatTensor<B>

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

fn softmax_native( tensor: FloatTensor<B>, dim: usize, logarithmic: bool, ) -> FloatTensor<B>

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

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

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

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

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

fn embedding(weights: FloatTensor<B>, indices: IntTensor<B>) -> FloatTensor<B>

Embedding operation.

Arguments

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

Returns

The output tensor.

fn embedding_backward( weights: FloatTensor<B>, output_grad: FloatTensor<B>, indices: IntTensor<B>, ) -> FloatTensor<B>

Embedding backward operation.

Arguments

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

Returns

The gradient.

fn linear( x: FloatTensor<B>, weight: FloatTensor<B>, bias: Option<FloatTensor<B>>, ) -> FloatTensor<B>

Linear transformation.

Shapes

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

fn linear_x_backward( weight: FloatTensor<B>, output_grad: FloatTensor<B>, ) -> FloatTensor<B>

Backward pass for linear, returning the gradient for x.

fn linear_weight_backward( x: FloatTensor<B>, output_grad: FloatTensor<B>, ) -> FloatTensor<B>

Backward pass for linear, returning the gradient for weight.

fn linear_bias_backward(output_grad: FloatTensor<B>) -> FloatTensor<B>

Backward pass for linear, returning the gradient for bias.

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

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: FloatTensor<B>, weight: FloatTensor<B>, output_grad: FloatTensor<B>, options: ConvOptions<1>, ) -> FloatTensor<B>

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

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

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

fn conv1d_bias_backward( x: FloatTensor<B>, bias: FloatTensor<B>, output_grad: FloatTensor<B>, ) -> FloatTensor<B>

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

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

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

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

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

fn conv2d_bias_backward( x: FloatTensor<B>, bias: FloatTensor<B>, output_grad: FloatTensor<B>, ) -> FloatTensor<B>

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

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

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

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

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

fn conv3d_bias_backward( x: FloatTensor<B>, bias: FloatTensor<B>, output_grad: FloatTensor<B>, ) -> FloatTensor<B>

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

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

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: FloatTensor<B>, output_grad: FloatTensor<B>, options: ConvTransposeOptions<1>, ) -> FloatTensor<B>

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

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

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

fn conv_transpose1d_bias_backward( x: FloatTensor<B>, bias: FloatTensor<B>, output_grad: FloatTensor<B>, ) -> FloatTensor<B>

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

fn conv_transpose2d_x_backward( weight: FloatTensor<B>, output_grad: FloatTensor<B>, options: ConvTransposeOptions<2>, ) -> FloatTensor<B>

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

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

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

fn conv_transpose2d_bias_backward( x: FloatTensor<B>, bias: FloatTensor<B>, output_grad: FloatTensor<B>, ) -> FloatTensor<B>

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

fn conv_transpose3d_x_backward( weight: FloatTensor<B>, output_grad: FloatTensor<B>, options: ConvTransposeOptions<3>, ) -> FloatTensor<B>

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

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

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

fn conv_transpose3d_bias_backward( x: FloatTensor<B>, bias: FloatTensor<B>, output_grad: FloatTensor<B>, ) -> FloatTensor<B>

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

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

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: FloatTensor<B>, kernel_size: usize, stride: usize, padding: usize, count_include_pad: bool, ceil_mode: bool, ) -> FloatTensor<B>

One dimensional avg pooling.

Shapes

x: [batch_size, channels, length],

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

Backward pass for the avg pooling 1d operation.

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

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: FloatTensor<B>, grad: FloatTensor<B>, kernel_size: [usize; 3], stride: [usize; 3], padding: [usize; 3], count_include_pad: bool, ceil_mode: bool, ) -> FloatTensor<B>

Input gradients for three dimensional average pooling.

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

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

fn adaptive_avg_pool3d_backward( x: FloatTensor<B>, grad: FloatTensor<B>, ) -> FloatTensor<B>

Input gradients for native three dimensional adaptive average pooling.

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

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

fn max_pool3d_with_indices( x: FloatTensor<B>, 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: FloatTensor<B>, grad: FloatTensor<B>, indices: IntTensor<B>, _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: FloatTensor<B>, output_size: usize) -> FloatTensor<B>

One dimensional adaptive avg pooling.

Shapes

x: [batch_size, channels, length],

fn adaptive_avg_pool1d_backward( x: FloatTensor<B>, grad: FloatTensor<B>, ) -> FloatTensor<B>

Backward pass for the adaptive avg pooling 1d operation.

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

One dimensional max pooling.

Shapes

x: [batch_size, channels, length],

fn max_pool1d_with_indices( x: FloatTensor<B>, 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: FloatTensor<B>, kernel_size: usize, stride: usize, padding: usize, dilation: usize, ceil_mode: bool, output_grad: FloatTensor<B>, indices: IntTensor<B>, ) -> MaxPool1dBackward<B>

Backward pass for the max pooling 1d operation.

fn interpolate1d( x: FloatTensor<B>, size: usize, options: InterpolateOptions, ) -> FloatTensor<B>

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

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

Backward line resizing in the original input geometry and storage.

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

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

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

Backward volume resizing through the original spatial and depth filters.

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

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: FloatTensor<B>, gamma: FloatTensor<B>, epsilon: f64, ) -> FloatTensor<B>

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

fn rms_norm_with_stats( tensor: FloatTensor<B>, gamma: FloatTensor<B>, epsilon: f64, ) -> RmsNormOutput<B>

RMSNorm forward retaining reciprocal row norms for backward.

fn rms_norm_backward( tensor: FloatTensor<B>, gamma: FloatTensor<B>, grad: FloatTensor<B>, rstd: FloatTensor<B>, ) -> RmsNormBackward<B>

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

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

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

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

Native forward with statistics retained for backward.

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

Native backward using the exact statistics returned by forward.

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

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

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

Differentiable primitive composition for backends without native LayerNorm backward.

fn ctc_loss( log_probs: FloatTensor<B>, targets: IntTensor<B>, input_lengths: IntTensor<B>, target_lengths: IntTensor<B>, blank: usize, ) -> FloatTensor<B>

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: FloatTensor<B>, _targets: IntTensor<B>, _input_lengths: IntTensor<B>, _target_lengths: IntTensor<B>, _grad_loss: FloatTensor<B>, _blank: usize, ) -> FloatTensor<B>

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