Trait burn_tensor::ops::ModuleOps

pub trait ModuleOps<B: Backend> {
    // Required methods
    fn conv2d(
        x: B::TensorPrimitive<4>,
        weight: B::TensorPrimitive<4>,
        bias: Option<B::TensorPrimitive<1>>,
        options: ConvOptions<2>
    ) -> B::TensorPrimitive<4>;
    fn conv_transpose2d(
        x: B::TensorPrimitive<4>,
        weight: B::TensorPrimitive<4>,
        bias: Option<B::TensorPrimitive<1>>,
        options: ConvTransposeOptions<2>
    ) -> B::TensorPrimitive<4>;
    fn avg_pool2d(
        x: B::TensorPrimitive<4>,
        kernel_size: [usize; 2],
        stride: [usize; 2],
        padding: [usize; 2]
    ) -> B::TensorPrimitive<4>;
    fn avg_pool2d_backward(
        x: B::TensorPrimitive<4>,
        grad: B::TensorPrimitive<4>,
        kernel_size: [usize; 2],
        stride: [usize; 2],
        padding: [usize; 2]
    ) -> B::TensorPrimitive<4>;
    fn max_pool2d(
        x: B::TensorPrimitive<4>,
        kernel_size: [usize; 2],
        stride: [usize; 2],
        padding: [usize; 2]
    ) -> B::TensorPrimitive<4>;
    fn max_pool2d_with_indices(
        x: B::TensorPrimitive<4>,
        kernel_size: [usize; 2],
        stride: [usize; 2],
        padding: [usize; 2]
    ) -> MaxPool2dWithIndices<B>;
    fn max_pool2d_with_indices_backward(
        x: B::TensorPrimitive<4>,
        kernel_size: [usize; 2],
        stride: [usize; 2],
        padding: [usize; 2],
        output_grad: B::TensorPrimitive<4>,
        indices: B::IntTensorPrimitive<4>
    ) -> MaxPool2dBackward<B>;

    // Provided methods
    fn embedding(
        weights: B::TensorPrimitive<2>,
        indices: B::IntTensorPrimitive<2>
    ) -> B::TensorPrimitive<3> { ... }
    fn embedding_backward(
        weights: B::TensorPrimitive<2>,
        output_grad: B::TensorPrimitive<3>,
        indices: B::IntTensorPrimitive<2>
    ) -> B::TensorPrimitive<2> { ... }
    fn conv2d_backward(
        x: B::TensorPrimitive<4>,
        weight: B::TensorPrimitive<4>,
        bias: Option<B::TensorPrimitive<1>>,
        output_grad: B::TensorPrimitive<4>,
        options: ConvOptions<2>
    ) -> Conv2dBackward<B> { ... }
    fn conv1d(
        x: B::TensorPrimitive<3>,
        weight: B::TensorPrimitive<3>,
        bias: Option<B::TensorPrimitive<1>>,
        options: ConvOptions<1>
    ) -> B::TensorPrimitive<3> { ... }
    fn conv_transpose1d(
        x: B::TensorPrimitive<3>,
        weight: B::TensorPrimitive<3>,
        bias: Option<B::TensorPrimitive<1>>,
        options: ConvTransposeOptions<1>
    ) -> B::TensorPrimitive<3> { ... }
    fn conv1d_backward(
        x: B::TensorPrimitive<3>,
        weight: B::TensorPrimitive<3>,
        bias: Option<B::TensorPrimitive<1>>,
        output_grad: B::TensorPrimitive<3>,
        options: ConvOptions<1>
    ) -> Conv1dBackward<B> { ... }
    fn avg_pool1d(
        x: B::TensorPrimitive<3>,
        kernel_size: usize,
        stride: usize,
        padding: usize
    ) -> B::TensorPrimitive<3> { ... }
    fn avg_pool1d_backward(
        x: B::TensorPrimitive<3>,
        grad: B::TensorPrimitive<3>,
        kernel_size: usize,
        stride: usize,
        padding: usize
    ) -> B::TensorPrimitive<3> { ... }
}

Module operations trait.

Required Methods

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

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 conv_transpose2d( x: B::TensorPrimitive<4>, weight: B::TensorPrimitive<4>, bias: Option<B::TensorPrimitive<1>>, options: ConvTransposeOptions<2> ) -> B::TensorPrimitive<4>

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 avg_pool2d( x: B::TensorPrimitive<4>, kernel_size: [usize; 2], stride: [usize; 2], padding: [usize; 2] ) -> B::TensorPrimitive<4>

Two dimensional avg pooling.

Shapes

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

fn avg_pool2d_backward( x: B::TensorPrimitive<4>, grad: B::TensorPrimitive<4>, kernel_size: [usize; 2], stride: [usize; 2], padding: [usize; 2] ) -> B::TensorPrimitive<4>

Backward pass for the avg pooling 2d operation.

fn max_pool2d( x: B::TensorPrimitive<4>, kernel_size: [usize; 2], stride: [usize; 2], padding: [usize; 2] ) -> B::TensorPrimitive<4>

Two dimensional max pooling.

Shapes

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

fn max_pool2d_with_indices( x: B::TensorPrimitive<4>, kernel_size: [usize; 2], stride: [usize; 2], padding: [usize; 2] ) -> MaxPool2dWithIndices<B>

Two dimensional max pooling with indices.

Shapes

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

fn max_pool2d_with_indices_backward( x: B::TensorPrimitive<4>, kernel_size: [usize; 2], stride: [usize; 2], padding: [usize; 2], output_grad: B::TensorPrimitive<4>, indices: B::IntTensorPrimitive<4> ) -> MaxPool2dBackward<B>

Backward pass for the max pooling 2d operation.

Provided Methods

fn embedding( weights: B::TensorPrimitive<2>, indices: B::IntTensorPrimitive<2> ) -> B::TensorPrimitive<3>

Embedding operation.

Arguments

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

Returns

The output tensor.

fn embedding_backward( weights: B::TensorPrimitive<2>, output_grad: B::TensorPrimitive<3>, indices: B::IntTensorPrimitive<2> ) -> B::TensorPrimitive<2>

Embedding backward operation.

Arguments

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

Returns

The gradient.

fn conv2d_backward( x: B::TensorPrimitive<4>, weight: B::TensorPrimitive<4>, bias: Option<B::TensorPrimitive<1>>, output_grad: B::TensorPrimitive<4>, options: ConvOptions<2> ) -> Conv2dBackward<B>

Backward pass for the conv2d operation.

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

One dimensional convolution.

Shapes

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

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

One dimensional transposed convolution.

Shapes

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

fn conv1d_backward( x: B::TensorPrimitive<3>, weight: B::TensorPrimitive<3>, bias: Option<B::TensorPrimitive<1>>, output_grad: B::TensorPrimitive<3>, options: ConvOptions<1> ) -> Conv1dBackward<B>

Backward pass for the conv1d operation.

fn avg_pool1d( x: B::TensorPrimitive<3>, kernel_size: usize, stride: usize, padding: usize ) -> B::TensorPrimitive<3>

One dimensional avg pooling.

Shapes

x: [batch_size, channels, length],

fn avg_pool1d_backward( x: B::TensorPrimitive<3>, grad: B::TensorPrimitive<3>, kernel_size: usize, stride: usize, padding: usize ) -> B::TensorPrimitive<3>

Backward pass for the avg pooling 1d operation.

Implementors