Trait ModuleOps 

pub trait ModuleOps<B>
where
    B: Backend,
{
    // Required methods
    fn conv2d(
        x: <B as Backend>::FloatTensorPrimitive,
        weight: <B as Backend>::FloatTensorPrimitive,
        bias: Option<<B as Backend>::FloatTensorPrimitive>,
        options: ConvOptions<2>,
    ) -> <B as Backend>::FloatTensorPrimitive;
    fn deform_conv2d(
        x: <B as Backend>::FloatTensorPrimitive,
        offset: <B as Backend>::FloatTensorPrimitive,
        weight: <B as Backend>::FloatTensorPrimitive,
        mask: Option<<B as Backend>::FloatTensorPrimitive>,
        bias: Option<<B as Backend>::FloatTensorPrimitive>,
        options: DeformConvOptions<2>,
    ) -> <B as Backend>::FloatTensorPrimitive;
    fn deform_conv2d_backward(
        x: <B as Backend>::FloatTensorPrimitive,
        offset: <B as Backend>::FloatTensorPrimitive,
        weight: <B as Backend>::FloatTensorPrimitive,
        mask: Option<<B as Backend>::FloatTensorPrimitive>,
        bias: Option<<B as Backend>::FloatTensorPrimitive>,
        output_grad: <B as Backend>::FloatTensorPrimitive,
        options: DeformConvOptions<2>,
    ) -> DeformConv2dBackward<B>;
    fn conv3d(
        x: <B as Backend>::FloatTensorPrimitive,
        weight: <B as Backend>::FloatTensorPrimitive,
        bias: Option<<B as Backend>::FloatTensorPrimitive>,
        options: ConvOptions<3>,
    ) -> <B as Backend>::FloatTensorPrimitive;
    fn conv_transpose2d(
        x: <B as Backend>::FloatTensorPrimitive,
        weight: <B as Backend>::FloatTensorPrimitive,
        bias: Option<<B as Backend>::FloatTensorPrimitive>,
        options: ConvTransposeOptions<2>,
    ) -> <B as Backend>::FloatTensorPrimitive;
    fn conv_transpose3d(
        x: <B as Backend>::FloatTensorPrimitive,
        weight: <B as Backend>::FloatTensorPrimitive,
        bias: Option<<B as Backend>::FloatTensorPrimitive>,
        options: ConvTransposeOptions<3>,
    ) -> <B as Backend>::FloatTensorPrimitive;
    fn avg_pool2d(
        x: <B as Backend>::FloatTensorPrimitive,
        kernel_size: [usize; 2],
        stride: [usize; 2],
        padding: [usize; 2],
        count_include_pad: bool,
        ceil_mode: bool,
    ) -> <B as Backend>::FloatTensorPrimitive;
    fn avg_pool2d_backward(
        x: <B as Backend>::FloatTensorPrimitive,
        grad: <B as Backend>::FloatTensorPrimitive,
        kernel_size: [usize; 2],
        stride: [usize; 2],
        padding: [usize; 2],
        count_include_pad: bool,
        ceil_mode: bool,
    ) -> <B as Backend>::FloatTensorPrimitive;
    fn adaptive_avg_pool2d(
        x: <B as Backend>::FloatTensorPrimitive,
        output_size: [usize; 2],
    ) -> <B as Backend>::FloatTensorPrimitive;
    fn adaptive_avg_pool2d_backward(
        x: <B as Backend>::FloatTensorPrimitive,
        grad: <B as Backend>::FloatTensorPrimitive,
    ) -> <B as Backend>::FloatTensorPrimitive;
    fn max_pool2d(
        x: <B as Backend>::FloatTensorPrimitive,
        kernel_size: [usize; 2],
        stride: [usize; 2],
        padding: [usize; 2],
        dilation: [usize; 2],
        ceil_mode: bool,
    ) -> <B as Backend>::FloatTensorPrimitive;
    fn max_pool2d_with_indices(
        x: <B as Backend>::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 Backend>::FloatTensorPrimitive,
        kernel_size: [usize; 2],
        stride: [usize; 2],
        padding: [usize; 2],
        dilation: [usize; 2],
        ceil_mode: bool,
        output_grad: <B as Backend>::FloatTensorPrimitive,
        indices: <B as Backend>::IntTensorPrimitive,
    ) -> MaxPool2dBackward<B>;
    fn interpolate(
        x: <B as Backend>::FloatTensorPrimitive,
        output_size: [usize; 2],
        options: InterpolateOptions,
    ) -> <B as Backend>::FloatTensorPrimitive;
    fn interpolate_backward(
        x: <B as Backend>::FloatTensorPrimitive,
        grad: <B as Backend>::FloatTensorPrimitive,
        output_size: [usize; 2],
        options: InterpolateOptions,
    ) -> <B as Backend>::FloatTensorPrimitive;

    // Provided methods
    fn embedding(
        weights: <B as Backend>::FloatTensorPrimitive,
        indices: <B as Backend>::IntTensorPrimitive,
    ) -> <B as Backend>::FloatTensorPrimitive { ... }
    fn embedding_backward(
        weights: <B as Backend>::FloatTensorPrimitive,
        output_grad: <B as Backend>::FloatTensorPrimitive,
        indices: <B as Backend>::IntTensorPrimitive,
    ) -> <B as Backend>::FloatTensorPrimitive { ... }
    fn conv1d(
        x: <B as Backend>::FloatTensorPrimitive,
        weight: <B as Backend>::FloatTensorPrimitive,
        bias: Option<<B as Backend>::FloatTensorPrimitive>,
        options: ConvOptions<1>,
    ) -> <B as Backend>::FloatTensorPrimitive { ... }
    fn conv1d_x_backward(
        x: <B as Backend>::FloatTensorPrimitive,
        weight: <B as Backend>::FloatTensorPrimitive,
        output_grad: <B as Backend>::FloatTensorPrimitive,
        options: ConvOptions<1>,
    ) -> <B as Backend>::FloatTensorPrimitive { ... }
    fn conv1d_weight_backward(
        x: <B as Backend>::FloatTensorPrimitive,
        weight: <B as Backend>::FloatTensorPrimitive,
        output_grad: <B as Backend>::FloatTensorPrimitive,
        options: ConvOptions<1>,
    ) -> <B as Backend>::FloatTensorPrimitive { ... }
    fn conv1d_bias_backward(
        x: <B as Backend>::FloatTensorPrimitive,
        bias: <B as Backend>::FloatTensorPrimitive,
        output_grad: <B as Backend>::FloatTensorPrimitive,
    ) -> <B as Backend>::FloatTensorPrimitive { ... }
    fn conv2d_x_backward(
        x: <B as Backend>::FloatTensorPrimitive,
        weight: <B as Backend>::FloatTensorPrimitive,
        output_grad: <B as Backend>::FloatTensorPrimitive,
        options: ConvOptions<2>,
    ) -> <B as Backend>::FloatTensorPrimitive { ... }
    fn conv2d_weight_backward(
        x: <B as Backend>::FloatTensorPrimitive,
        weight: <B as Backend>::FloatTensorPrimitive,
        output_grad: <B as Backend>::FloatTensorPrimitive,
        options: ConvOptions<2>,
    ) -> <B as Backend>::FloatTensorPrimitive { ... }
    fn conv2d_bias_backward(
        x: <B as Backend>::FloatTensorPrimitive,
        weight: <B as Backend>::FloatTensorPrimitive,
        bias: <B as Backend>::FloatTensorPrimitive,
        output_grad: <B as Backend>::FloatTensorPrimitive,
    ) -> <B as Backend>::FloatTensorPrimitive { ... }
    fn conv3d_x_backward(
        x: <B as Backend>::FloatTensorPrimitive,
        weight: <B as Backend>::FloatTensorPrimitive,
        output_grad: <B as Backend>::FloatTensorPrimitive,
        options: ConvOptions<3>,
    ) -> <B as Backend>::FloatTensorPrimitive { ... }
    fn conv3d_weight_backward(
        x: <B as Backend>::FloatTensorPrimitive,
        weight: <B as Backend>::FloatTensorPrimitive,
        output_grad: <B as Backend>::FloatTensorPrimitive,
        options: ConvOptions<3>,
    ) -> <B as Backend>::FloatTensorPrimitive { ... }
    fn conv3d_bias_backward(
        x: <B as Backend>::FloatTensorPrimitive,
        weight: <B as Backend>::FloatTensorPrimitive,
        bias: <B as Backend>::FloatTensorPrimitive,
        output_grad: <B as Backend>::FloatTensorPrimitive,
    ) -> <B as Backend>::FloatTensorPrimitive { ... }
    fn conv_transpose1d(
        x: <B as Backend>::FloatTensorPrimitive,
        weight: <B as Backend>::FloatTensorPrimitive,
        bias: Option<<B as Backend>::FloatTensorPrimitive>,
        options: ConvTransposeOptions<1>,
    ) -> <B as Backend>::FloatTensorPrimitive { ... }
    fn conv_transpose1d_x_backward(
        weight: <B as Backend>::FloatTensorPrimitive,
        output_grad: <B as Backend>::FloatTensorPrimitive,
        options: ConvTransposeOptions<1>,
    ) -> <B as Backend>::FloatTensorPrimitive { ... }
    fn conv_transpose1d_weight_backward(
        x: <B as Backend>::FloatTensorPrimitive,
        weight: <B as Backend>::FloatTensorPrimitive,
        output_grad: <B as Backend>::FloatTensorPrimitive,
        options: ConvTransposeOptions<1>,
    ) -> <B as Backend>::FloatTensorPrimitive { ... }
    fn conv_transpose1d_bias_backward(
        x: <B as Backend>::FloatTensorPrimitive,
        bias: <B as Backend>::FloatTensorPrimitive,
        output_grad: <B as Backend>::FloatTensorPrimitive,
    ) -> <B as Backend>::FloatTensorPrimitive { ... }
    fn conv_transpose2d_x_backward(
        weight: <B as Backend>::FloatTensorPrimitive,
        output_grad: <B as Backend>::FloatTensorPrimitive,
        options: ConvTransposeOptions<2>,
    ) -> <B as Backend>::FloatTensorPrimitive { ... }
    fn conv_transpose2d_weight_backward(
        x: <B as Backend>::FloatTensorPrimitive,
        weight: <B as Backend>::FloatTensorPrimitive,
        output_grad: <B as Backend>::FloatTensorPrimitive,
        options: ConvTransposeOptions<2>,
    ) -> <B as Backend>::FloatTensorPrimitive { ... }
    fn conv_transpose2d_bias_backward(
        x: <B as Backend>::FloatTensorPrimitive,
        bias: <B as Backend>::FloatTensorPrimitive,
        output_grad: <B as Backend>::FloatTensorPrimitive,
    ) -> <B as Backend>::FloatTensorPrimitive { ... }
    fn conv_transpose3d_x_backward(
        weight: <B as Backend>::FloatTensorPrimitive,
        output_grad: <B as Backend>::FloatTensorPrimitive,
        options: ConvTransposeOptions<3>,
    ) -> <B as Backend>::FloatTensorPrimitive { ... }
    fn conv_transpose3d_weight_backward(
        x: <B as Backend>::FloatTensorPrimitive,
        weight: <B as Backend>::FloatTensorPrimitive,
        output_grad: <B as Backend>::FloatTensorPrimitive,
        options: ConvTransposeOptions<3>,
    ) -> <B as Backend>::FloatTensorPrimitive { ... }
    fn conv_transpose3d_bias_backward(
        x: <B as Backend>::FloatTensorPrimitive,
        bias: <B as Backend>::FloatTensorPrimitive,
        output_grad: <B as Backend>::FloatTensorPrimitive,
    ) -> <B as Backend>::FloatTensorPrimitive { ... }
    fn unfold4d(
        x: <B as Backend>::FloatTensorPrimitive,
        kernel_size: [usize; 2],
        options: UnfoldOptions,
    ) -> <B as Backend>::FloatTensorPrimitive { ... }
    fn avg_pool1d(
        x: <B as Backend>::FloatTensorPrimitive,
        kernel_size: usize,
        stride: usize,
        padding: usize,
        count_include_pad: bool,
        ceil_mode: bool,
    ) -> <B as Backend>::FloatTensorPrimitive { ... }
    fn avg_pool1d_backward(
        x: <B as Backend>::FloatTensorPrimitive,
        grad: <B as Backend>::FloatTensorPrimitive,
        kernel_size: usize,
        stride: usize,
        padding: usize,
        count_include_pad: bool,
        ceil_mode: bool,
    ) -> <B as Backend>::FloatTensorPrimitive { ... }
    fn adaptive_avg_pool1d(
        x: <B as Backend>::FloatTensorPrimitive,
        output_size: usize,
    ) -> <B as Backend>::FloatTensorPrimitive { ... }
    fn adaptive_avg_pool1d_backward(
        x: <B as Backend>::FloatTensorPrimitive,
        grad: <B as Backend>::FloatTensorPrimitive,
    ) -> <B as Backend>::FloatTensorPrimitive { ... }
    fn max_pool1d(
        x: <B as Backend>::FloatTensorPrimitive,
        kernel_size: usize,
        stride: usize,
        padding: usize,
        dilation: usize,
        ceil_mode: bool,
    ) -> <B as Backend>::FloatTensorPrimitive { ... }
    fn max_pool1d_with_indices(
        x: <B as Backend>::FloatTensorPrimitive,
        kernel_size: usize,
        stride: usize,
        padding: usize,
        dilation: usize,
        ceil_mode: bool,
    ) -> MaxPool1dWithIndices<B> { ... }
    fn max_pool1d_with_indices_backward(
        x: <B as Backend>::FloatTensorPrimitive,
        kernel_size: usize,
        stride: usize,
        padding: usize,
        dilation: usize,
        ceil_mode: bool,
        output_grad: <B as Backend>::FloatTensorPrimitive,
        indices: <B as Backend>::IntTensorPrimitive,
    ) -> MaxPool1dBackward<B> { ... }
    fn attention(
        query: <B as Backend>::FloatTensorPrimitive,
        key: <B as Backend>::FloatTensorPrimitive,
        value: <B as Backend>::FloatTensorPrimitive,
        mask: Option<<B as Backend>::BoolTensorPrimitive>,
    ) -> <B as Backend>::FloatTensorPrimitive { ... }
}

Module operations trait.

Required Methods

fn conv2d( x: <B as Backend>::FloatTensorPrimitive, weight: <B as Backend>::FloatTensorPrimitive, bias: Option<<B as Backend>::FloatTensorPrimitive>, options: ConvOptions<2>, ) -> <B as Backend>::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 Backend>::FloatTensorPrimitive, offset: <B as Backend>::FloatTensorPrimitive, weight: <B as Backend>::FloatTensorPrimitive, mask: Option<<B as Backend>::FloatTensorPrimitive>, bias: Option<<B as Backend>::FloatTensorPrimitive>, options: DeformConvOptions<2>, ) -> <B as Backend>::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 Backend>::FloatTensorPrimitive, offset: <B as Backend>::FloatTensorPrimitive, weight: <B as Backend>::FloatTensorPrimitive, mask: Option<<B as Backend>::FloatTensorPrimitive>, bias: Option<<B as Backend>::FloatTensorPrimitive>, output_grad: <B as Backend>::FloatTensorPrimitive, options: DeformConvOptions<2>, ) -> DeformConv2dBackward<B>

Backward pass for the deform_conv2d operation.

fn conv3d( x: <B as Backend>::FloatTensorPrimitive, weight: <B as Backend>::FloatTensorPrimitive, bias: Option<<B as Backend>::FloatTensorPrimitive>, options: ConvOptions<3>, ) -> <B as Backend>::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 Backend>::FloatTensorPrimitive, weight: <B as Backend>::FloatTensorPrimitive, bias: Option<<B as Backend>::FloatTensorPrimitive>, options: ConvTransposeOptions<2>, ) -> <B as Backend>::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 Backend>::FloatTensorPrimitive, weight: <B as Backend>::FloatTensorPrimitive, bias: Option<<B as Backend>::FloatTensorPrimitive>, options: ConvTransposeOptions<3>, ) -> <B as Backend>::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 Backend>::FloatTensorPrimitive, kernel_size: [usize; 2], stride: [usize; 2], padding: [usize; 2], count_include_pad: bool, ceil_mode: bool, ) -> <B as Backend>::FloatTensorPrimitive

Two dimensional avg pooling.

Shapes

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

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

Backward pass for the avg pooling 2d operation.

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

Two dimensional adaptive avg pooling.

Shapes

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

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

Backward pass for the adaptive avg pooling 2d operation.

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

Two dimensional max pooling.

Shapes

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

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

Backward pass for the max pooling 2d operation.

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

Down/up samples the input.

Shapes

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

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

Backward pass for the interpolate operation.

Provided Methods

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

Embedding operation.

Arguments

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

Returns

The output tensor.

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

Embedding backward operation.

Arguments

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

Returns

The gradient.

fn conv1d( x: <B as Backend>::FloatTensorPrimitive, weight: <B as Backend>::FloatTensorPrimitive, bias: Option<<B as Backend>::FloatTensorPrimitive>, options: ConvOptions<1>, ) -> <B as Backend>::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 Backend>::FloatTensorPrimitive, weight: <B as Backend>::FloatTensorPrimitive, output_grad: <B as Backend>::FloatTensorPrimitive, options: ConvOptions<1>, ) -> <B as Backend>::FloatTensorPrimitive

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

fn conv_transpose1d( x: <B as Backend>::FloatTensorPrimitive, weight: <B as Backend>::FloatTensorPrimitive, bias: Option<<B as Backend>::FloatTensorPrimitive>, options: ConvTransposeOptions<1>, ) -> <B as Backend>::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 Backend>::FloatTensorPrimitive, output_grad: <B as Backend>::FloatTensorPrimitive, options: ConvTransposeOptions<1>, ) -> <B as Backend>::FloatTensorPrimitive

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

fn unfold4d( x: <B as Backend>::FloatTensorPrimitive, kernel_size: [usize; 2], options: UnfoldOptions, ) -> <B as Backend>::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 Backend>::FloatTensorPrimitive, kernel_size: usize, stride: usize, padding: usize, count_include_pad: bool, ceil_mode: bool, ) -> <B as Backend>::FloatTensorPrimitive

One dimensional avg pooling.

Shapes

x: [batch_size, channels, length],

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

Backward pass for the avg pooling 1d operation.

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

One dimensional adaptive avg pooling.

Shapes

x: [batch_size, channels, length],

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

Backward pass for the adaptive avg pooling 1d operation.

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

One dimensional max pooling.

Shapes

x: [batch_size, channels, length],

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

Backward pass for the max pooling 1d operation.

fn attention( query: <B as Backend>::FloatTensorPrimitive, key: <B as Backend>::FloatTensorPrimitive, value: <B as Backend>::FloatTensorPrimitive, mask: Option<<B as Backend>::BoolTensorPrimitive>, ) -> <B as Backend>::FloatTensorPrimitive

Computes scaled dot-product attention: softmax(QKᵗ / √d) · V, optionally applying a mask 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], here true indicates positions to mask (i.e. set to -∞ before softmax).

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.

Dyn Compatibility

This trait is not dyn compatible.

In older versions of Rust, dyn compatibility was called "object safety", so this trait is not object safe.

Implementations on Foreign Types

impl<E, I, Q> ModuleOps<NdArray<E, I, Q>> for NdArray<E, I, Q>

where E: FloatNdArrayElement, I: IntNdArrayElement, Q: QuantElement, NdArrayTensor: From<ArrayBase<OwnedArcRepr<E>, Dim<IxDynImpl>>> + From<ArrayBase<OwnedArcRepr<I>, Dim<IxDynImpl>>>,

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

fn deform_conv2d( x: <NdArray<E, I, Q> as Backend>::FloatTensorPrimitive, offset: <NdArray<E, I, Q> as Backend>::FloatTensorPrimitive, weight: <NdArray<E, I, Q> as Backend>::FloatTensorPrimitive, mask: Option<<NdArray<E, I, Q> as Backend>::FloatTensorPrimitive>, bias: Option<<NdArray<E, I, Q> as Backend>::FloatTensorPrimitive>, options: DeformConvOptions<2>, ) -> <NdArray<E, I, Q> as Backend>::FloatTensorPrimitive

fn deform_conv2d_backward( x: <NdArray<E, I, Q> as Backend>::FloatTensorPrimitive, offset: <NdArray<E, I, Q> as Backend>::FloatTensorPrimitive, weight: <NdArray<E, I, Q> as Backend>::FloatTensorPrimitive, mask: Option<<NdArray<E, I, Q> as Backend>::FloatTensorPrimitive>, bias: Option<<NdArray<E, I, Q> as Backend>::FloatTensorPrimitive>, output_grad: <NdArray<E, I, Q> as Backend>::FloatTensorPrimitive, options: DeformConvOptions<2>, ) -> DeformConv2dBackward<NdArray<E, I, Q>>

fn conv_transpose2d( x: <NdArray<E, I, Q> as Backend>::FloatTensorPrimitive, weight: <NdArray<E, I, Q> as Backend>::FloatTensorPrimitive, bias: Option<<NdArray<E, I, Q> as Backend>::FloatTensorPrimitive>, options: ConvTransposeOptions<2>, ) -> <NdArray<E, I, Q> as Backend>::FloatTensorPrimitive

fn avg_pool2d( x: <NdArray<E, I, Q> as Backend>::FloatTensorPrimitive, kernel_size: [usize; 2], stride: [usize; 2], padding: [usize; 2], count_include_pad: bool, ceil_mode: bool, ) -> <NdArray<E, I, Q> as Backend>::FloatTensorPrimitive

fn avg_pool2d_backward( x: <NdArray<E, I, Q> as Backend>::FloatTensorPrimitive, grad: <NdArray<E, I, Q> as Backend>::FloatTensorPrimitive, kernel_size: [usize; 2], stride: [usize; 2], padding: [usize; 2], count_include_pad: bool, ceil_mode: bool, ) -> <NdArray<E, I, Q> as Backend>::FloatTensorPrimitive

fn max_pool2d( x: <NdArray<E, I, Q> as Backend>::FloatTensorPrimitive, kernel_size: [usize; 2], stride: [usize; 2], padding: [usize; 2], dilation: [usize; 2], ceil_mode: bool, ) -> <NdArray<E, I, Q> as Backend>::FloatTensorPrimitive

fn max_pool2d_with_indices( x: <NdArray<E, I, Q> as Backend>::FloatTensorPrimitive, kernel_size: [usize; 2], stride: [usize; 2], padding: [usize; 2], dilation: [usize; 2], ceil_mode: bool, ) -> MaxPool2dWithIndices<NdArray<E, I, Q>>

fn max_pool2d_with_indices_backward( x: <NdArray<E, I, Q> as Backend>::FloatTensorPrimitive, kernel_size: [usize; 2], stride: [usize; 2], padding: [usize; 2], dilation: [usize; 2], ceil_mode: bool, output_grad: <NdArray<E, I, Q> as Backend>::FloatTensorPrimitive, indices: NdArrayTensor, ) -> MaxPool2dBackward<NdArray<E, I, Q>>

fn adaptive_avg_pool2d( x: <NdArray<E, I, Q> as Backend>::FloatTensorPrimitive, output_size: [usize; 2], ) -> <NdArray<E, I, Q> as Backend>::FloatTensorPrimitive

fn adaptive_avg_pool2d_backward( x: <NdArray<E, I, Q> as Backend>::FloatTensorPrimitive, grad: <NdArray<E, I, Q> as Backend>::FloatTensorPrimitive, ) -> <NdArray<E, I, Q> as Backend>::FloatTensorPrimitive

fn interpolate( x: <NdArray<E, I, Q> as Backend>::FloatTensorPrimitive, output_size: [usize; 2], options: InterpolateOptions, ) -> <NdArray<E, I, Q> as Backend>::FloatTensorPrimitive

fn interpolate_backward( x: <NdArray<E, I, Q> as Backend>::FloatTensorPrimitive, grad: <NdArray<E, I, Q> as Backend>::FloatTensorPrimitive, output_size: [usize; 2], options: InterpolateOptions, ) -> <NdArray<E, I, Q> as Backend>::FloatTensorPrimitive

fn conv3d( x: <NdArray<E, I, Q> as Backend>::FloatTensorPrimitive, weight: <NdArray<E, I, Q> as Backend>::FloatTensorPrimitive, bias: Option<<NdArray<E, I, Q> as Backend>::FloatTensorPrimitive>, options: ConvOptions<3>, ) -> <NdArray<E, I, Q> as Backend>::FloatTensorPrimitive

fn conv_transpose3d( x: <NdArray<E, I, Q> as Backend>::FloatTensorPrimitive, weight: <NdArray<E, I, Q> as Backend>::FloatTensorPrimitive, bias: Option<<NdArray<E, I, Q> as Backend>::FloatTensorPrimitive>, options: ConvTransposeOptions<3>, ) -> <NdArray<E, I, Q> as Backend>::FloatTensorPrimitive

Implementors

impl<B, C> ModuleOps<Autodiff<B, C>> for Autodiff<B, C>

where B: Backend, C: CheckpointStrategy,