burn_tensor/tensor/
module.rs

1use burn_backend::ops::ModuleOps;
2use burn_dispatch::Dispatch;
3use burn_std::{MatmulTransformAction, MatmulTransformAnalysis, MatmulTransformPolicy};
4
5use crate::{
6    Bool, DType, Int, Tensor, check,
7    check::TensorCheck,
8    ops::{
9        AttentionModuleOptions, BridgeTensor, ConvOptions, ConvTransposeOptions, DeformConvOptions,
10        InterpolateOptions, PadMode, PaddedConvOptions, UnfoldOptions,
11    },
12};
13
14/// Applies batch normalization using explicitly supplied channel statistics.
15///
16/// `input` has shape `[batch, channels, ...]`; `gamma`, `beta`, `mean`, and
17/// `variance` each have shape `[channels]`.
18///
19/// This function doesn't calculate or update statistics. Callers may supply
20/// running statistics for inference or batch statistics calculated by a
21/// training path.
22pub fn batch_norm<const D: usize>(
23    input: Tensor<D>,
24    gamma: Tensor<1>,
25    beta: Tensor<1>,
26    mean: Tensor<1>,
27    variance: Tensor<1>,
28    epsilon: f64,
29) -> Tensor<D> {
30    assert!(D >= 2, "batch norm requires an input rank of at least 2");
31    let channels = input.dims()[1];
32    assert_eq!(gamma.dims(), [channels], "invalid batch norm gamma shape");
33    assert_eq!(beta.dims(), [channels], "invalid batch norm beta shape");
34    assert_eq!(mean.dims(), [channels], "invalid batch norm mean shape");
35    assert_eq!(
36        variance.dims(),
37        [channels],
38        "invalid batch norm variance shape"
39    );
40    Tensor::new(BridgeTensor::float(Dispatch::batch_norm(
41        input.primitive.into_float(),
42        gamma.primitive.into_float(),
43        beta.primitive.into_float(),
44        mean.primitive.into_float(),
45        variance.primitive.into_float(),
46        epsilon,
47    )))
48}
49
50/// Computes the [CTC loss](burn_backend::ops::ModuleOps::ctc_loss).
51///
52/// # Arguments
53///
54/// * `log_probs` - Log-probabilities of shape `[T, N, C]`
55/// * `targets` - Target label indices of shape `[N, S]`
56/// * `input_lengths` - Actual input sequence lengths per batch element `[N]`
57/// * `target_lengths` - Actual target lengths per batch element `[N]`
58/// * `blank` - Index of the blank label
59///
60/// # Returns
61///
62/// Per-sample loss of shape `[N]`
63pub fn ctc_loss(
64    log_probs: Tensor<3>,
65    targets: Tensor<2, Int>,
66    input_lengths: Tensor<1, Int>,
67    target_lengths: Tensor<1, Int>,
68    blank: usize,
69) -> Tensor<1> {
70    Tensor::new(BridgeTensor::float(Dispatch::ctc_loss(
71        log_probs.primitive.into_float(),
72        targets.primitive.into(),
73        input_lengths.primitive.into(),
74        target_lengths.primitive.into(),
75        blank,
76    )))
77}
78
79/// Applies the [embedding module](burn_backend::ops::ModuleOps::embedding).
80pub fn embedding(weights: Tensor<2>, indices: Tensor<2, Int>) -> Tensor<3> {
81    Tensor::new(BridgeTensor::float(Dispatch::embedding(
82        weights.primitive.into_float(),
83        indices.primitive.into(),
84    )))
85}
86
87/// Applies a [1D convolution](burn_backend::ops::ModuleOps::conv1d).
88///
89/// Accepts [`ConvOptions`] for symmetric padding, or [`PaddedConvOptions`] for
90/// asymmetric padding. When asymmetric padding is specified, an explicit pad
91/// operation is applied before the convolution backend op.
92pub fn conv1d(
93    x: Tensor<3>,
94    weight: Tensor<3>,
95    bias: Option<Tensor<1>>,
96    options: impl Into<PaddedConvOptions<1>>,
97) -> Tensor<3> {
98    let padded_options = options.into();
99    check!(TensorCheck::conv(
100        "conv1d",
101        x.dims(),
102        weight.dims(),
103        padded_options.options.groups,
104    ));
105
106    if let Some(padding_end) = padded_options.padding_end {
107        let left = padded_options.options.padding[0];
108        let right = padding_end[0];
109        // For 1D (NCL format), pad the length dimension
110        let padded = x.pad((left, right, 0, 0), PadMode::Constant(0.0));
111        let zero_options = ConvOptions::new(
112            padded_options.options.stride,
113            [0],
114            padded_options.options.dilation,
115            padded_options.options.groups,
116        );
117        Tensor::new(BridgeTensor::float(Dispatch::conv1d(
118            padded.primitive.into_float(),
119            weight.primitive.into_float(),
120            bias.map(|b| b.primitive.into_float()),
121            zero_options,
122        )))
123    } else {
124        Tensor::new(BridgeTensor::float(Dispatch::conv1d(
125            x.primitive.into_float(),
126            weight.primitive.into_float(),
127            bias.map(|b| b.primitive.into_float()),
128            padded_options.options,
129        )))
130    }
131}
132
133/// Applies a [2D convolution](burn_backend::ops::ModuleOps::conv2d).
134///
135/// Accepts [`ConvOptions`] for symmetric padding, or [`PaddedConvOptions`] for
136/// asymmetric padding. When asymmetric padding is specified, an explicit pad
137/// operation is applied before the convolution backend op.
138pub fn conv2d(
139    x: Tensor<4>,
140    weight: Tensor<4>,
141    bias: Option<Tensor<1>>,
142    options: impl Into<PaddedConvOptions<2>>,
143) -> Tensor<4> {
144    let padded_options = options.into();
145    check!(TensorCheck::conv(
146        "conv2d",
147        x.dims(),
148        weight.dims(),
149        padded_options.options.groups,
150    ));
151
152    if let Some(padding_end) = padded_options.padding_end {
153        let top = padded_options.options.padding[0];
154        let left = padded_options.options.padding[1];
155        let bottom = padding_end[0];
156        let right = padding_end[1];
157        // For 2D (NCHW format), pad height and width
158        let padded = x.pad((left, right, top, bottom), PadMode::Constant(0.0));
159        let zero_options = ConvOptions::new(
160            padded_options.options.stride,
161            [0, 0],
162            padded_options.options.dilation,
163            padded_options.options.groups,
164        );
165        Tensor::new(BridgeTensor::float(Dispatch::conv2d(
166            padded.primitive.into_float(),
167            weight.primitive.into_float(),
168            bias.map(|b| b.primitive.into_float()),
169            zero_options,
170        )))
171    } else {
172        Tensor::new(BridgeTensor::float(Dispatch::conv2d(
173            x.primitive.into_float(),
174            weight.primitive.into_float(),
175            bias.map(|b| b.primitive.into_float()),
176            padded_options.options,
177        )))
178    }
179}
180
181/// Applies a [3D convolution](burn_backend::ops::ModuleOps::conv3d).
182///
183/// Accepts [`ConvOptions`] for symmetric padding, or [`PaddedConvOptions`] for
184/// asymmetric padding. Asymmetric 3D padding is not yet supported.
185pub fn conv3d(
186    x: Tensor<5>,
187    weight: Tensor<5>,
188    bias: Option<Tensor<1>>,
189    options: impl Into<PaddedConvOptions<3>>,
190) -> Tensor<5> {
191    let padded_options = options.into();
192    check!(TensorCheck::conv(
193        "conv3d",
194        x.dims(),
195        weight.dims(),
196        padded_options.options.groups,
197    ));
198
199    if padded_options.is_asymmetric() {
200        panic!("Asymmetric padding is not yet supported for conv3d");
201    }
202
203    Tensor::new(BridgeTensor::float(Dispatch::conv3d(
204        x.primitive.into_float(),
205        weight.primitive.into_float(),
206        bias.map(|b| b.primitive.into_float()),
207        padded_options.options,
208    )))
209}
210
211/// Applies a [Deformable 2D convolution](burn_backend::ops::ModuleOps::deform_conv2d).
212pub fn deform_conv2d(
213    x: Tensor<4>,
214    offset: Tensor<4>,
215    weight: Tensor<4>,
216    mask: Option<Tensor<4>>,
217    bias: Option<Tensor<1>>,
218    options: DeformConvOptions<2>,
219) -> Tensor<4> {
220    check!(TensorCheck::conv(
221        "deform_conv2d",
222        x.dims(),
223        weight.dims(),
224        options.weight_groups,
225    ));
226    Tensor::new(BridgeTensor::float(Dispatch::deform_conv2d(
227        x.primitive.into_float(),
228        offset.primitive.into_float(),
229        weight.primitive.into_float(),
230        mask.map(|m| m.primitive.into_float()),
231        bias.map(|b| b.primitive.into_float()),
232        options,
233    )))
234}
235
236/// Applies a [1D transposed convolution](burn_backend::ops::ModuleOps::conv_transpose1d).
237pub fn conv_transpose1d(
238    x: Tensor<3>,
239    weight: Tensor<3>,
240    bias: Option<Tensor<1>>,
241    options: ConvTransposeOptions<1>,
242) -> Tensor<3> {
243    check!(TensorCheck::conv_transpose(
244        "conv_transpose1d",
245        x.dims(),
246        weight.dims(),
247    ));
248    Tensor::new(BridgeTensor::float(Dispatch::conv_transpose1d(
249        x.primitive.into_float(),
250        weight.primitive.into_float(),
251        bias.map(|b| b.primitive.into_float()),
252        options,
253    )))
254}
255
256/// Applies a [2D transposed convolution](burn_backend::ops::ModuleOps::conv_transpose2d).
257pub fn conv_transpose2d(
258    x: Tensor<4>,
259    weight: Tensor<4>,
260    bias: Option<Tensor<1>>,
261    options: ConvTransposeOptions<2>,
262) -> Tensor<4> {
263    check!(TensorCheck::conv_transpose(
264        "conv_transpose2d",
265        x.dims(),
266        weight.dims(),
267    ));
268    Tensor::new(BridgeTensor::float(Dispatch::conv_transpose2d(
269        x.primitive.into_float(),
270        weight.primitive.into_float(),
271        bias.map(|b| b.primitive.into_float()),
272        options,
273    )))
274}
275
276/// Applies a 3D transposed convolution](burn_backend::ops::ModuleOps::conv_transpose3d).
277pub fn conv_transpose3d(
278    x: Tensor<5>,
279    weight: Tensor<5>,
280    bias: Option<Tensor<1>>,
281    options: ConvTransposeOptions<3>,
282) -> Tensor<5> {
283    check!(TensorCheck::conv_transpose(
284        "conv_transpose3d",
285        x.dims(),
286        weight.dims(),
287    ));
288    Tensor::new(BridgeTensor::float(Dispatch::conv_transpose3d(
289        x.primitive.into_float(),
290        weight.primitive.into_float(),
291        bias.map(|b| b.primitive.into_float()),
292        options,
293    )))
294}
295
296/// Applies a [4D to 3D unfold](burn_backend::ops::ModuleOps::unfold4d).
297pub fn unfold4d(x: Tensor<4>, kernel_size: [usize; 2], options: UnfoldOptions) -> Tensor<3> {
298    Tensor::new(BridgeTensor::float(Dispatch::unfold4d(
299        x.primitive.into_float(),
300        kernel_size,
301        options,
302    )))
303}
304
305/// Applies a 3D to 4D fold, the inverse of [unfold4d].
306///
307/// Combines an array of sliding local blocks into a large containing tensor, summing the
308/// values of blocks that overlap. This is the operation performed by
309/// [`torch.nn.Fold`](https://pytorch.org/docs/stable/generated/torch.nn.Fold.html), and is the
310/// adjoint of [unfold4d]: it reuses the same one-hot kernel through a [conv_transpose2d].
311///
312/// # Arguments
313///
314/// * `x` - Input columns of shape
315///   `[batch_size, channels * kernel_size_0 * kernel_size_1, number_of_blocks]`.
316/// * `output_size` - The spatial size `[height, width]` of the folded output tensor.
317/// * `kernel_size` - The size of the sliding blocks.
318/// * `options` - The stride, padding and dilation of the matching unfold.
319///
320/// # Returns
321///
322/// A tensor of shape `[batch_size, channels, output_size_0, output_size_1]`.
323pub fn fold4d(
324    x: Tensor<3>,
325    output_size: [usize; 2],
326    kernel_size: [usize; 2],
327    options: UnfoldOptions,
328) -> Tensor<4> {
329    Tensor::new(BridgeTensor::float(Dispatch::fold4d(
330        x.primitive.into_float(),
331        output_size,
332        kernel_size,
333        options,
334    )))
335}
336
337/// Applies a [1D max pooling](burn_backend::ops::ModuleOps::max_pool1d).
338pub fn max_pool1d(
339    x: Tensor<3>,
340    kernel_size: usize,
341    stride: usize,
342    padding: usize,
343    dilation: usize,
344    ceil_mode: bool,
345) -> Tensor<3> {
346    Tensor::new(BridgeTensor::float(Dispatch::max_pool1d(
347        x.primitive.into_float(),
348        kernel_size,
349        stride,
350        padding,
351        dilation,
352        ceil_mode,
353    )))
354}
355
356/// Applies a [2D max pooling](burn_backend::ops::ModuleOps::max_pool2d).
357pub fn max_pool2d(
358    x: Tensor<4>,
359    kernel_size: [usize; 2],
360    stride: [usize; 2],
361    padding: [usize; 2],
362    dilation: [usize; 2],
363    ceil_mode: bool,
364) -> Tensor<4> {
365    Tensor::new(BridgeTensor::float(Dispatch::max_pool2d(
366        x.primitive.into_float(),
367        kernel_size,
368        stride,
369        padding,
370        dilation,
371        ceil_mode,
372    )))
373}
374
375/// Applies a [2D avg pooling](burn_backend::ops::ModuleOps::avg_pool2d).
376pub fn avg_pool2d(
377    x: Tensor<4>,
378    kernel_size: [usize; 2],
379    stride: [usize; 2],
380    padding: [usize; 2],
381    count_include_pad: bool,
382    ceil_mode: bool,
383) -> Tensor<4> {
384    Tensor::new(BridgeTensor::float(Dispatch::avg_pool2d(
385        x.primitive.into_float(),
386        kernel_size,
387        stride,
388        padding,
389        count_include_pad,
390        ceil_mode,
391    )))
392}
393
394/// Applies a [1D avg pooling](burn_backend::ops::ModuleOps::avg_pool1d).
395pub fn avg_pool1d(
396    x: Tensor<3>,
397    kernel_size: usize,
398    stride: usize,
399    padding: usize,
400    count_include_pad: bool,
401    ceil_mode: bool,
402) -> Tensor<3> {
403    Tensor::new(BridgeTensor::float(Dispatch::avg_pool1d(
404        x.primitive.into_float(),
405        kernel_size,
406        stride,
407        padding,
408        count_include_pad,
409        ceil_mode,
410    )))
411}
412
413/// Applies a [1D max pooling](burn_backend::ops::ModuleOps::max_pool1d).
414pub fn max_pool1d_with_indices(
415    x: Tensor<3>,
416    kernel_size: usize,
417    stride: usize,
418    padding: usize,
419    dilation: usize,
420    ceil_mode: bool,
421) -> (Tensor<3>, Tensor<3, Int>) {
422    let indices_dtype = x.device().settings().int_dtype;
423    let output = Dispatch::max_pool1d_with_indices(
424        x.primitive.into_float(),
425        kernel_size,
426        stride,
427        padding,
428        dilation,
429        ceil_mode,
430        indices_dtype,
431    );
432
433    (
434        Tensor::new(BridgeTensor::float(output.output)),
435        Tensor::new(BridgeTensor::int(output.indices)),
436    )
437}
438
439/// Applies a [2D max pooling with indices](burn_backend::ops::ModuleOps::max_pool2d_with_indices).
440pub fn max_pool2d_with_indices(
441    x: Tensor<4>,
442    kernel_size: [usize; 2],
443    stride: [usize; 2],
444    padding: [usize; 2],
445    dilation: [usize; 2],
446    ceil_mode: bool,
447) -> (Tensor<4>, Tensor<4, Int>) {
448    let indices_dtype = x.device().settings().int_dtype;
449    let output = Dispatch::max_pool2d_with_indices(
450        x.primitive.into_float(),
451        kernel_size,
452        stride,
453        padding,
454        dilation,
455        ceil_mode,
456        indices_dtype,
457    );
458
459    (
460        Tensor::new(BridgeTensor::float(output.output)),
461        Tensor::new(BridgeTensor::int(output.indices)),
462    )
463}
464
465/// Applies a [2D adaptive avg pooling](burn_backend::ops::ModuleOps::adaptive_avg_pool2d).
466pub fn adaptive_avg_pool2d(x: Tensor<4>, output_size: [usize; 2]) -> Tensor<4> {
467    Tensor::new(BridgeTensor::float(Dispatch::adaptive_avg_pool2d(
468        x.primitive.into_float(),
469        output_size,
470    )))
471}
472
473/// Applies a [3D adaptive avg pooling](burn_backend::ops::ModuleOps::adaptive_avg_pool3d).
474pub fn adaptive_avg_pool3d(x: Tensor<5>, output_size: [usize; 3]) -> Tensor<5> {
475    Tensor::new(BridgeTensor::float(Dispatch::adaptive_avg_pool3d(
476        x.primitive.into_float(),
477        output_size,
478    )))
479}
480
481/// Applies a [1D adaptive avg pooling](burn_backend::ops::ModuleOps::adaptive_avg_pool1d).
482pub fn adaptive_avg_pool1d(x: Tensor<3>, output_size: usize) -> Tensor<3> {
483    Tensor::new(BridgeTensor::float(Dispatch::adaptive_avg_pool1d(
484        x.primitive.into_float(),
485        output_size,
486    )))
487}
488
489/// Applies a [2D interpolation](burn_backend::ops::ModuleOps::interpolate).
490pub fn interpolate(
491    x: Tensor<4>,
492    output_size: [usize; 2],
493    options: InterpolateOptions,
494) -> Tensor<4> {
495    Tensor::new(BridgeTensor::float(Dispatch::interpolate(
496        x.primitive.into_float(),
497        output_size,
498        options,
499    )))
500}
501
502/// Applies a linear transformation to the input tensor using the given weight and bias.
503///
504/// ```math
505/// y = x @ weight + [bias]
506/// ```
507///
508/// # Arguments:
509///
510/// - `input` is the input tensor, ``[..., d_input]``.
511/// - `weight` is the weight tensor, ``[d_input, d_output]``.
512/// - `bias` is the bias tensor (optional), ``[d_output]``.
513///
514/// # Returns:
515///
516/// The transformed tensor, ``[..., d_output]``.
517///
518/// # Compatibility
519///
520/// This function differs from PyTorch's ``torch.nn.functional.linear`` in that it does not
521/// transpose the weight matrix. In PyTorch, the weight matrix is transposed before
522/// multiplication:
523///
524/// ```math
525/// y = x @ weight^T + [bias]
526/// ```
527pub fn linear<const D: usize>(
528    input: Tensor<D>,
529    weight: Tensor<2>,
530    bias: Option<Tensor<1>>,
531) -> Tensor<D> {
532    if D == 1 {
533        // Insert and remove an extra batch dimension for the batch matmul to work.
534        let input = input.unsqueeze::<2>();
535        let output = linear(input, weight, bias);
536        return output.squeeze_dim(0);
537    }
538
539    // A quantized weight must stay quantized: `linear_impl` converts its
540    // operands to float, which would dequantize (materialize) the whole weight
541    // matrix on every forward. Route through the quantized matmul instead, which
542    // streams the packed weight directly — but reuse the same batch-fold policy
543    // the float `linear` applies, so a decode-shaped call folds its batches into
544    // the rows for one `[rows, d_in] @ [d_in, d_out]` matmul rather than a
545    // broadcast batched matmul that re-reads the packed weight per batch.
546    if let DType::QFloat(_) = weight.dtype() {
547        let dims = input.dims();
548        let analysis = MatmulTransformAnalysis::from_shapes(&input.shape(), &weight.shape());
549
550        let output = match MatmulTransformPolicy::default().action(&analysis) {
551            MatmulTransformAction::MergeBatches { rows } => {
552                let d_in = dims[D - 1];
553                let d_out = weight.dims()[1];
554
555                let folded = input.reshape([rows, d_in]).matmul(weight);
556
557                let mut out_dims = dims;
558                out_dims[D - 1] = d_out;
559                folded.reshape(out_dims)
560            }
561            MatmulTransformAction::Keep => input.matmul(weight.unsqueeze::<D>()),
562        };
563
564        return match bias {
565            Some(bias) => output + bias.unsqueeze(),
566            None => output,
567        };
568    }
569
570    Tensor::new(linear_impl(
571        input.primitive,
572        weight.primitive,
573        bias.map(|b| b.primitive),
574    ))
575}
576
577fn linear_impl(
578    input: BridgeTensor,
579    weight: BridgeTensor,
580    bias: Option<BridgeTensor>,
581) -> BridgeTensor {
582    BridgeTensor::float(Dispatch::linear(
583        input.into_float(),
584        weight.into_float(),
585        bias.map(|b| b.into_float()),
586    ))
587}
588
589/// Computes scaled dot-product attention: softmax(QKᵗ * scale) · V,
590/// where scale defaults to 1/sqrt(head_dim) (configurable via `options.scale`).
591/// Optionally applies masking, additive bias, causal masking, and softcap.
592///
593/// # Arguments
594/// - `query`: Query tensor of shape `[batch_size, num_heads, seq_len_q, head_dim]`
595/// - `key`: Key tensor of shape `[batch_size, num_heads, seq_len_k, head_dim]`
596/// - `value`: Value tensor of shape `[batch_size, num_heads, seq_len_k, val_dim]`
597/// - `mask`: Optional boolean mask of shape `[batch_size, num_heads, seq_len_q, seq_len_k]`,
598///   where `true` indicates positions to mask (i.e. set to -inf before softmax).
599/// - `attn_bias`: Optional float tensor of shape `[batch_size, num_heads, seq_len_q, seq_len_k]`
600///   added to the attention scores before softmax (e.g. ALiBi, relative position biases).
601/// - `options`: Additional attention options (custom scale, softcap, causal masking).
602///
603/// # Returns
604/// A tensor of shape `[batch_size, num_heads, seq_len_q, val_dim]`
605/// representing the attended context per head.
606///
607/// # Note
608/// This implementation does not support dropout and is intended for inference or
609/// use cases where dropout is not needed.
610pub fn attention(
611    query: Tensor<4>,
612    key: Tensor<4>,
613    value: Tensor<4>,
614    mask: Option<Tensor<4, Bool>>,
615    attn_bias: Option<Tensor<4>>,
616    options: AttentionModuleOptions,
617) -> Tensor<4> {
618    Tensor::new(BridgeTensor::float(Dispatch::attention(
619        query.primitive.into_float(),
620        key.primitive.into_float(),
621        value.primitive.into_float(),
622        mask.map(|mask| mask.primitive.into()),
623        attn_bias.map(|bias| bias.primitive.into_float()),
624        options,
625    )))
626}
627
628/// Exports attention fallback to test backend's attention against.
629pub fn attention_fallback(
630    query: Tensor<4>,
631    key: Tensor<4>,
632    value: Tensor<4>,
633    mask: Option<Tensor<4, Bool>>,
634    attn_bias: Option<Tensor<4>>,
635    options: AttentionModuleOptions,
636) -> Tensor<4> {
637    Tensor::new(BridgeTensor::float(
638        burn_backend::ops::attention::attention_fallback::<Dispatch>(
639            query.primitive.into_float(),
640            key.primitive.into_float(),
641            value.primitive.into_float(),
642            mask.map(|mask| mask.primitive.into()),
643            attn_bias.map(|bias| bias.primitive.into_float()),
644            options,
645        ),
646    ))
647}
648
649/// Calculate the [2D convolution](burn_backend::ops::ModuleOps::conv2d) backward pass, returning the gradient for `weight`.
650pub fn conv2d_weight_backward(
651    x: Tensor<4>,
652    weight: Tensor<4>,
653    output_grad: Tensor<4>,
654    options: ConvOptions<2>,
655) -> Tensor<4> {
656    Tensor::new(BridgeTensor::float(Dispatch::conv2d_weight_backward(
657        x.primitive.into_float(),
658        weight.primitive.into_float(),
659        output_grad.primitive.into_float(),
660        options,
661    )))
662}
663
664/// Backward pass for the [avg pooling 2d](ModuleOps::avg_pool2d) operation.
665pub fn avg_pool2d_backward(
666    x: Tensor<4>,
667    grad: Tensor<4>,
668    kernel_size: [usize; 2],
669    stride: [usize; 2],
670    padding: [usize; 2],
671    count_include_pad: bool,
672    ceil_mode: bool,
673) -> Tensor<4> {
674    Tensor::new(BridgeTensor::float(Dispatch::avg_pool2d_backward(
675        x.primitive.into_float(),
676        grad.primitive.into_float(),
677        kernel_size,
678        stride,
679        padding,
680        count_include_pad,
681        ceil_mode,
682    )))
683}
684
685/// Backward pass for the [max pooling 2d](ModuleOps::max_pool2d_with_indices) operation.
686#[allow(clippy::too_many_arguments)]
687pub fn max_pool2d_with_indices_backward(
688    x: Tensor<4>,
689    kernel_size: [usize; 2],
690    stride: [usize; 2],
691    padding: [usize; 2],
692    dilation: [usize; 2],
693    ceil_mode: bool,
694    output_grad: Tensor<4>,
695    indices: Tensor<4, Int>,
696) -> Tensor<4> {
697    Tensor::new(BridgeTensor::float(
698        Dispatch::max_pool2d_with_indices_backward(
699            x.primitive.into_float(),
700            kernel_size,
701            stride,
702            padding,
703            dilation,
704            ceil_mode,
705            output_grad.primitive.into_float(),
706            indices.primitive.into(),
707        )
708        .x_grad,
709    ))
710}
711
712/// Applies Layer Normalization over the last dimension of the input tensor.
713///
714/// Computes `(x - mean) / sqrt(var + epsilon) * gamma + beta`, where `mean` and
715/// (biased) `var` are reduced over the last axis.
716///
717/// # Shapes
718///
719/// - input: `[..., any, d_model]`
720/// - output: `[..., any, d_model]`
721pub fn layer_norm<const D: usize>(
722    input: Tensor<D>,
723    gamma: Tensor<1>,
724    beta: Option<Tensor<1>>,
725    epsilon: f64,
726) -> Tensor<D> {
727    Tensor::new(layer_norm_impl(
728        input.primitive,
729        gamma.primitive,
730        beta.map(|b| b.primitive),
731        epsilon,
732    ))
733}
734
735fn layer_norm_impl(
736    input: BridgeTensor,
737    gamma: BridgeTensor,
738    beta: Option<BridgeTensor>,
739    epsilon: f64,
740) -> BridgeTensor {
741    BridgeTensor::float(Dispatch::layer_norm(
742        input.into_float(),
743        gamma.into_float(),
744        beta.map(|b| b.into_float()),
745        epsilon,
746    ))
747}
burn_tensor/tensor/module.rs

burn_tensor/tensor/
module.rs
