Available on crate feature
api only.Expand description
The neural network module.
Functions§
- adaptive_
avg_ pool1d - Applies a 1D adaptive avg pooling.
- adaptive_
avg_ pool2d - Applies a 2D adaptive avg pooling.
- adaptive_
avg_ pool3d - Adaptive average pooling to explicit
[depth, height, width]extents. - attention
- Computes scaled dot-product attention: softmax(QKᵗ * scale) · V,
where scale defaults to 1/sqrt(head_dim) (configurable via
options.scale). Optionally applies masking, additive bias, causal masking, and softcap. - attention_
fallback - Exports attention fallback to test backend’s attention against.
- avg_
pool1d - Applies a 1D avg pooling.
- avg_
pool1d_ padded - Average pooling with explicit
(left, right)padding. - avg_
pool2d - Applies a 2D avg pooling.
- avg_
pool2d_ padded - Average pooling with explicit height and width
(before, after)pairs. - avg_
pool3d - Average-pool native
[batch, channels, depth, height, width]activations. - avg_
pool3d_ padded - Average pooling with depth, height and width
(before, after)pairs. - conv1d
- Applies a 1D convolution.
- conv2d
- Applies a 2D convolution.
- conv3d
- Applies a 3D convolution.
- conv_
transpose1d - Applies a 1D transposed convolution.
- conv_
transpose1d_ with_ output_ size - Apply a 1D transposed convolution with an explicit output length.
- conv_
transpose2d - Applies a 2D transposed convolution.
- conv_
transpose2d_ with_ output_ size - Apply a 2D transposed convolution with explicit output height and width.
- conv_
transpose3d - Applies a 3D transposed convolution](crate::api::ops::ModuleOps::conv_transpose3d).
- conv_
transpose3d_ with_ output_ size - Apply a 3D transposed convolution with explicit output depth, height and width.
- ctc_
loss - Computes the CTC loss.
- deform_
conv2d - Applies a Deformable 2D convolution.
- embedding
- Applies the embedding module.
- group_
norm - Channel GroupNorm on
[batch, channels, ...], with independently optional affine leaves. Uses biased variance and saved FP32 statistics for low-precision training; F64 retains FP64. - interpolate
- Applies a 2D interpolation.
- interpolate1d
- Resize native
[batch, channels, length]activations. - interpolate3d
- Resize native
[batch, channels, depth, height, width]activations. - linear
- Applies a linear transformation to the input tensor using the given weight and bias.
- max_
pool1d - Applies a 1D max pooling.
- max_
pool1d_ padded - Maximum pooling with explicit
(left, right)padding and native dilation. - max_
pool1d_ with_ indices - Applies a 1D max pooling.
- max_
pool1d_ with_ indices_ padded - Maximum pooling with asymmetric length padding and original input positions.
- max_
pool2d - Applies a 2D max pooling.
- max_
pool2d_ padded - Maximum pooling with explicit height and width
(before, after)pairs. - max_
pool2d_ with_ indices - Applies a 2D max pooling with indices.
- max_
pool2d_ with_ indices_ padded - Maximum pooling with asymmetric spatial padding and original
h * W + windices. - max_
pool3d - Maximum-pool native volumes using the backend’s volume operation.
- max_
pool3d_ padded - Maximum pooling with explicit depth, height and width padding pairs.
- max_
pool3d_ with_ indices - Native volume maxima and flattened I64 positions in the original input.
Empty selections are
-1; tie/NaN selection retains spatial-then-depth order. - max_
pool3d_ with_ indices_ padded - Maximum volume pooling with asymmetric padding and original flattened positions.
- rms_
norm - Last-axis RMSNorm with working arithmetic and one output-storage cast. F32/F16/BF16 activations use FP32 statistics; F64 activations retain FP64.
- unfold4d
- Applies a 4D to 3D unfold.