burn_backend/backend/ops/modules/base.rs
1use super::{conv, ctc, linear, pool};
2use crate::ops::unfold::{create_unfolding_weight, unfold4d_using_conv2d};
3use crate::tensor::{BoolTensor, FloatTensor, IntTensor};
4use crate::{Backend, TensorMetadata};
5pub use burn_std::ops::{
6 AttentionModuleOptions, ConvOptions, ConvTransposeOptions, DeformConvOptions,
7 GridSampleOptions, GridSamplePaddingMode, InterpolateMode, InterpolateOptions, PadMode,
8 PaddedConvOptions, UnfoldOptions,
9};
10use burn_std::{IntDType, Shape};
11
12/// Gradient computed during the backward pass for each tensor used by [conv2d](ModuleOps::conv2d).
13#[derive(new)]
14pub struct Conv2dBackward<B: Backend> {
15 /// Gradient.
16 pub x_grad: FloatTensor<B>,
17
18 /// Weights gradient.
19 pub weights_grad: FloatTensor<B>,
20
21 /// Bias gradient.
22 pub bias_grad: Option<FloatTensor<B>>,
23}
24
25/// Gradient computed during the backward pass for each tensor used by [deform_conv2d](ModuleOps::deform_conv2d).
26#[derive(new)]
27pub struct DeformConv2dBackward<B: Backend> {
28 /// Gradient.
29 pub x_grad: FloatTensor<B>,
30
31 /// Offset gradient.
32 pub offset_grad: FloatTensor<B>,
33
34 /// Weights gradient.
35 pub weight_grad: FloatTensor<B>,
36
37 /// Mask gradient.
38 pub mask_grad: Option<FloatTensor<B>>,
39
40 /// Bias gradient.
41 pub bias_grad: Option<FloatTensor<B>>,
42}
43
44/// Gradient computed during the backward pass for each tensor used by [conv3d](ModuleOps::conv3d).
45#[derive(new)]
46pub struct Conv3dBackward<B: Backend> {
47 /// Gradient.
48 pub x_grad: FloatTensor<B>,
49
50 /// Weights gradient.
51 pub weights_grad: FloatTensor<B>,
52
53 /// Bias gradient.
54 pub bias_grad: Option<FloatTensor<B>>,
55}
56
57/// Gradient computed during the backward pass for each tensor used by [max_pool1d](ModuleOps::max_pool1d).
58#[derive(new)]
59pub struct MaxPool1dBackward<B: Backend> {
60 /// Gradient.
61 pub x_grad: FloatTensor<B>,
62}
63
64/// Results from [max_pool1d](ModuleOps::max_pool1d_with_indices).
65#[derive(new)]
66pub struct MaxPool1dWithIndices<B: Backend> {
67 /// The output tensor.
68 pub output: FloatTensor<B>,
69
70 /// The indices tensor.
71 pub indices: IntTensor<B>,
72}
73
74/// Gradient computed during the backward pass for each tensor used by [max_pool2d](ModuleOps::max_pool2d).
75#[derive(new)]
76pub struct MaxPool2dBackward<B: Backend> {
77 /// Gradient.
78 pub x_grad: FloatTensor<B>,
79}
80
81/// Results from [max_pool2d](ModuleOps::max_pool2d_with_indices).
82#[derive(new)]
83pub struct MaxPool2dWithIndices<B: Backend> {
84 /// The output tensor.
85 pub output: FloatTensor<B>,
86
87 /// The indices tensor.
88 pub indices: IntTensor<B>,
89}
90
91/// Gradient computed during the backward pass for each tensor used by [interpolate](ModuleOps::interpolate).
92#[derive(new)]
93pub struct InterpolateBackward<B: Backend> {
94 /// Gradient.
95 pub x_grad: FloatTensor<B>,
96}
97
98/// Module operations trait.
99pub trait ModuleOps<B: Backend> {
100 /// Embedding operation.
101 ///
102 /// # Arguments
103 ///
104 /// * `weights` - The embedding weights.
105 /// * `indices` - The indices tensor.
106 ///
107 /// # Returns
108 ///
109 /// The output tensor.
110 fn embedding(weights: FloatTensor<B>, indices: IntTensor<B>) -> FloatTensor<B> {
111 let [batch_size, seq_length] = indices.shape().dims();
112 let [_, d_model] = weights.shape().dims();
113
114 let indices = B::int_reshape(indices, Shape::new([batch_size * seq_length]));
115 let output = B::float_select(weights, 0, indices);
116
117 B::float_reshape(output, Shape::new([batch_size, seq_length, d_model]))
118 }
119
120 /// Embedding backward operation.
121 ///
122 /// # Arguments
123 ///
124 /// * `weights` - The embedding weights.
125 /// * `output_grad` - The output gradient.
126 /// * `indices` - The indices tensor.
127 ///
128 /// # Returns
129 ///
130 /// The gradient.
131 fn embedding_backward(
132 weights: FloatTensor<B>,
133 output_grad: FloatTensor<B>,
134 indices: IntTensor<B>,
135 ) -> FloatTensor<B> {
136 let [batch_size, seq_length] = indices.shape().dims();
137 let [n_embeddings, d_model] = weights.shape().dims();
138 let device = weights.device();
139 let dtype = output_grad.dtype();
140
141 let indices = B::int_reshape(indices, Shape::new([batch_size * seq_length]));
142 let output_grad =
143 B::float_reshape(output_grad, Shape::new([batch_size * seq_length, d_model]));
144 let grad = B::float_zeros(Shape::new([n_embeddings, d_model]), &device, dtype.into());
145
146 B::float_select_add(grad, 0, indices, output_grad)
147 }
148
149 /// Linear transformation.
150 ///
151 /// # Shapes
152 ///
153 /// x: `[..., d_input]`,
154 /// weight: `[d_input, d_output]`,
155 /// bias: `[d_output]`,
156 fn linear(
157 x: FloatTensor<B>,
158 weight: FloatTensor<B>,
159 bias: Option<FloatTensor<B>>,
160 ) -> FloatTensor<B> {
161 linear::linear::<B>(x, weight, bias)
162 }
163 /// Backward pass for [linear](ModuleOps::linear), returning the gradient for `x`.
164 fn linear_x_backward(weight: FloatTensor<B>, output_grad: FloatTensor<B>) -> FloatTensor<B> {
165 linear::linear_x_backward::<B>(weight, output_grad)
166 }
167 /// Backward pass for [linear](ModuleOps::linear), returning the gradient for `weight`.
168 fn linear_weight_backward(x: FloatTensor<B>, output_grad: FloatTensor<B>) -> FloatTensor<B> {
169 linear::linear_weight_backward::<B>(x, output_grad)
170 }
171 /// Backward pass for [linear](ModuleOps::linear), returning the gradient for `bias`.
172 fn linear_bias_backward(output_grad: FloatTensor<B>) -> FloatTensor<B> {
173 linear::linear_bias_backward::<B>(output_grad)
174 }
175
176 /// One dimensional convolution.
177 ///
178 /// # Shapes
179 ///
180 /// x: `[batch_size, channels_in, length]`,
181 /// weight: `[channels_out, channels_in, kernel_size]`,
182 /// bias: `[channels_out]`,
183 fn conv1d(
184 x: FloatTensor<B>,
185 weight: FloatTensor<B>,
186 bias: Option<FloatTensor<B>>,
187 options: ConvOptions<1>,
188 ) -> FloatTensor<B> {
189 conv::conv1d_from_conv2d::<B>(x, weight, bias, options)
190 }
191 /// Backward pass for the [conv1d](ModuleOps::conv1d) operation, returning the gradient for `x`.
192 fn conv1d_x_backward(
193 x: FloatTensor<B>,
194 weight: FloatTensor<B>,
195 output_grad: FloatTensor<B>,
196 options: ConvOptions<1>,
197 ) -> FloatTensor<B> {
198 conv::conv1d_x_backward::<B>(x, weight, output_grad, options)
199 }
200 /// Backward pass for the [conv1d](ModuleOps::conv1d) operation, returning the gradient for `weight`.
201 fn conv1d_weight_backward(
202 x: FloatTensor<B>,
203 weight: FloatTensor<B>,
204 output_grad: FloatTensor<B>,
205 options: ConvOptions<1>,
206 ) -> FloatTensor<B> {
207 conv::conv1d_weight_backward::<B>(x, weight, output_grad, options)
208 }
209 /// Backward pass for the [conv1d](ModuleOps::conv1d) operation, returning the gradient for `bias`.
210 fn conv1d_bias_backward(
211 x: FloatTensor<B>,
212 bias: FloatTensor<B>,
213 output_grad: FloatTensor<B>,
214 ) -> FloatTensor<B> {
215 conv::conv1d_bias_backward::<B>(x, bias, output_grad)
216 }
217 /// Two dimensional convolution.
218 ///
219 /// # Shapes
220 ///
221 /// x: `[batch_size, channels_in, height, width]`,
222 /// weight: `[channels_out, channels_in, kernel_size_1, kernel_size_2]`,
223 /// bias: `[channels_out]`,
224 fn conv2d(
225 x: FloatTensor<B>,
226 weight: FloatTensor<B>,
227 bias: Option<FloatTensor<B>>,
228 options: ConvOptions<2>,
229 ) -> FloatTensor<B>;
230 /// Backward pass for the [conv2d](ModuleOps::conv2d) operation, returning the gradient for `x`.
231 fn conv2d_x_backward(
232 x: FloatTensor<B>,
233 weight: FloatTensor<B>,
234 output_grad: FloatTensor<B>,
235 options: ConvOptions<2>,
236 ) -> FloatTensor<B> {
237 conv::conv2d_x_backward::<B>(x, weight, output_grad, options)
238 }
239 /// Backward pass for the [conv2d](ModuleOps::conv2d) operation, returning the gradient for `weight`.
240 fn conv2d_weight_backward(
241 x: FloatTensor<B>,
242 weight: FloatTensor<B>,
243 output_grad: FloatTensor<B>,
244 options: ConvOptions<2>,
245 ) -> FloatTensor<B> {
246 conv::conv2d_weight_backward::<B>(x, weight, output_grad, options)
247 }
248 /// Backward pass for the [conv2d](ModuleOps::conv2d) operation, returning the gradient for `bias`.
249 fn conv2d_bias_backward(
250 x: FloatTensor<B>,
251 bias: FloatTensor<B>,
252 output_grad: FloatTensor<B>,
253 ) -> FloatTensor<B> {
254 conv::conv2d_bias_backward::<B>(x, bias, output_grad)
255 }
256
257 /// Two dimensional deformable convolution.
258 ///
259 /// # Shapes
260 ///
261 /// x: `[batch_size, channels_in, height, width]`,
262 /// weight: `[channels_out, channels_in, kernel_size_1, kernel_size_2]`,
263 /// bias: `[channels_out]`,
264 fn deform_conv2d(
265 x: FloatTensor<B>,
266 offset: FloatTensor<B>,
267 weight: FloatTensor<B>,
268 mask: Option<FloatTensor<B>>,
269 bias: Option<FloatTensor<B>>,
270 options: DeformConvOptions<2>,
271 ) -> FloatTensor<B>;
272 /// Backward pass for the [deform_conv2d](ModuleOps::deform_conv2d) operation.
273 fn deform_conv2d_backward(
274 x: FloatTensor<B>,
275 offset: FloatTensor<B>,
276 weight: FloatTensor<B>,
277 mask: Option<FloatTensor<B>>,
278 bias: Option<FloatTensor<B>>,
279 output_grad: FloatTensor<B>,
280 options: DeformConvOptions<2>,
281 ) -> DeformConv2dBackward<B>;
282
283 /// Three dimensional convolution.
284 ///
285 /// # Shapes
286 ///
287 /// x: `[batch_size, channels_in, depth, height, width]`,
288 /// weight: `[channels_out, channels_in, kernel_size_1, kernel_size_2, kernel_size_3]`,
289 /// bias: `[channels_out]`,
290 fn conv3d(
291 x: FloatTensor<B>,
292 weight: FloatTensor<B>,
293 bias: Option<FloatTensor<B>>,
294 options: ConvOptions<3>,
295 ) -> FloatTensor<B>;
296 /// Backward pass for the [conv3d](ModuleOps::conv3d) operation, returning the gradient for `x`.
297 fn conv3d_x_backward(
298 x: FloatTensor<B>,
299 weight: FloatTensor<B>,
300 output_grad: FloatTensor<B>,
301 options: ConvOptions<3>,
302 ) -> FloatTensor<B> {
303 conv::conv3d_x_backward::<B>(x, weight, output_grad, options)
304 }
305 /// Backward pass for the [conv3d](ModuleOps::conv3d) operation, returning the gradient for `weight`.
306 fn conv3d_weight_backward(
307 x: FloatTensor<B>,
308 weight: FloatTensor<B>,
309 output_grad: FloatTensor<B>,
310 options: ConvOptions<3>,
311 ) -> FloatTensor<B> {
312 conv::conv3d_weight_backward::<B>(x, weight, output_grad, options)
313 }
314 /// Backward pass for the [conv3d](ModuleOps::conv3d) operation, returning the gradient for `bias`.
315 fn conv3d_bias_backward(
316 x: FloatTensor<B>,
317 bias: FloatTensor<B>,
318 output_grad: FloatTensor<B>,
319 ) -> FloatTensor<B> {
320 conv::conv3d_bias_backward::<B>(x, bias, output_grad)
321 }
322 /// One dimensional transposed convolution.
323 ///
324 /// # Shapes
325 ///
326 /// x: `[batch_size, channels_in, length]`,
327 /// weight: `[channels_in, channels_out, length]`,
328 /// bias: `[channels_out]`,
329 fn conv_transpose1d(
330 x: FloatTensor<B>,
331 weight: FloatTensor<B>,
332 bias: Option<FloatTensor<B>>,
333 options: ConvTransposeOptions<1>,
334 ) -> FloatTensor<B> {
335 conv::conv_transpose1d_from_conv_transpose2d::<B>(x, weight, bias, options)
336 }
337 /// Backward pass for the [conv transpose 1d](ModuleOps::conv_transpose1d) operation, returning the gradient for `x`.
338 fn conv_transpose1d_x_backward(
339 weight: FloatTensor<B>,
340 output_grad: FloatTensor<B>,
341 options: ConvTransposeOptions<1>,
342 ) -> FloatTensor<B> {
343 conv::conv_transpose1d_x_backward::<B>(weight, output_grad, options)
344 }
345 /// Backward pass for the [conv transpose 1d](ModuleOps::conv_transpose1d) operation, returning the gradient for `weight`.
346 fn conv_transpose1d_weight_backward(
347 x: FloatTensor<B>,
348 weight: FloatTensor<B>,
349 output_grad: FloatTensor<B>,
350 options: ConvTransposeOptions<1>,
351 ) -> FloatTensor<B> {
352 conv::conv_transpose1d_weight_backward::<B>(x, weight, output_grad, options)
353 }
354 /// Backward pass for the [conv transpose 1d](ModuleOps::conv_transpose1d) operation, returning the gradient for `bias`.
355 fn conv_transpose1d_bias_backward(
356 x: FloatTensor<B>,
357 bias: FloatTensor<B>,
358 output_grad: FloatTensor<B>,
359 ) -> FloatTensor<B> {
360 conv::conv_transpose1d_bias_backward::<B>(x, bias, output_grad)
361 }
362
363 /// Two dimensional transposed convolution.
364 ///
365 /// # Shapes
366 ///
367 /// x: `[batch_size, channels_in, height, width]`,
368 /// weight: `[channels_in, channels_out, kernel_size_1, kernel_size_2]`,
369 /// bias: `[channels_out]`,
370 fn conv_transpose2d(
371 x: FloatTensor<B>,
372 weight: FloatTensor<B>,
373 bias: Option<FloatTensor<B>>,
374 options: ConvTransposeOptions<2>,
375 ) -> FloatTensor<B>;
376 /// Backward pass for the [conv transpose 2d](ModuleOps::conv_transpose2d) operation, returning the gradient for `x`.
377 fn conv_transpose2d_x_backward(
378 weight: FloatTensor<B>,
379 output_grad: FloatTensor<B>,
380 options: ConvTransposeOptions<2>,
381 ) -> FloatTensor<B> {
382 conv::conv_transpose2d_x_backward::<B>(weight, output_grad, options)
383 }
384 /// Backward pass for the [conv transpose 2d](ModuleOps::conv_transpose2d) operation, returning the gradient for `weight`.
385 fn conv_transpose2d_weight_backward(
386 x: FloatTensor<B>,
387 weight: FloatTensor<B>,
388 output_grad: FloatTensor<B>,
389 options: ConvTransposeOptions<2>,
390 ) -> FloatTensor<B> {
391 conv::conv_transpose2d_weight_backward::<B>(x, weight, output_grad, options)
392 }
393 /// Backward pass for the [conv transpose 2d](ModuleOps::conv_transpose2d) operation, returning the gradient for `bias`.
394 fn conv_transpose2d_bias_backward(
395 x: FloatTensor<B>,
396 bias: FloatTensor<B>,
397 output_grad: FloatTensor<B>,
398 ) -> FloatTensor<B> {
399 conv::conv_transpose2d_bias_backward::<B>(x, bias, output_grad)
400 }
401
402 /// Three dimensional transposed convolution.
403 ///
404 /// # Shapes
405 ///
406 /// x: `[batch_size, channels_in, height, width]`,
407 /// weight: `[channels_in, channels_out, kernel_size_1, kernel_size_2, kernel_size_3]`,
408 /// bias: `[channels_out]`,
409 fn conv_transpose3d(
410 x: FloatTensor<B>,
411 weight: FloatTensor<B>,
412 bias: Option<FloatTensor<B>>,
413 options: ConvTransposeOptions<3>,
414 ) -> FloatTensor<B>;
415 /// Backward pass for the [conv transpose 3d](ModuleOps::conv_transpose3d) operation, returning the gradient for `x`.
416 fn conv_transpose3d_x_backward(
417 weight: FloatTensor<B>,
418 output_grad: FloatTensor<B>,
419 options: ConvTransposeOptions<3>,
420 ) -> FloatTensor<B> {
421 conv::conv_transpose3d_x_backward::<B>(weight, output_grad, options)
422 }
423 /// Backward pass for the [conv transpose 3d](ModuleOps::conv_transpose3d) operation, returning the gradient for `weight`.
424 fn conv_transpose3d_weight_backward(
425 x: FloatTensor<B>,
426 weight: FloatTensor<B>,
427 output_grad: FloatTensor<B>,
428 options: ConvTransposeOptions<3>,
429 ) -> FloatTensor<B> {
430 conv::conv_transpose3d_weight_backward::<B>(x, weight, output_grad, options)
431 }
432 /// Backward pass for the [conv transpose 3d](ModuleOps::conv_transpose3d) operation, returning the gradient for `bias`.
433 fn conv_transpose3d_bias_backward(
434 x: FloatTensor<B>,
435 bias: FloatTensor<B>,
436 output_grad: FloatTensor<B>,
437 ) -> FloatTensor<B> {
438 conv::conv_transpose3d_bias_backward::<B>(x, bias, output_grad)
439 }
440
441 /// Four-dimensional unfolding.
442 ///
443 /// # Shapes
444 ///
445 /// * x: ``[batch_size, channels_in, height, width]``,
446 /// * returns: ``[batch_size, channels_in * kernel_size_1 * kernel_size_2, number of blocks]``,
447 fn unfold4d(
448 x: FloatTensor<B>,
449 kernel_size: [usize; 2],
450 options: UnfoldOptions,
451 ) -> FloatTensor<B> {
452 if options.padding == [0, 0] && options.dilation == [1, 1] {
453 let blocks = B::float_unfold(x, 2, kernel_size[0], options.stride[0]);
454 let blocks = B::float_unfold(blocks, 3, kernel_size[1], options.stride[1]);
455
456 // batch, channels, h_blocks, w_blocks, h_kern, w_kern
457
458 let blocks = B::float_permute(blocks, &[0, 1, 4, 5, 2, 3]);
459 let shape = blocks.shape();
460
461 // batch, channels, h_kern, w_kern, h_blocks, w_blocks
462
463 B::float_reshape(
464 blocks,
465 [
466 shape[0],
467 shape[1] * shape[2] * shape[3],
468 shape[4] * shape[5],
469 ]
470 .into(),
471 )
472 } else {
473 unfold4d_using_conv2d::<B>(x, kernel_size, options)
474 }
475 }
476
477 /// Four dimensional fold (`col2im`), the adjoint of [unfold4d](ModuleOps::unfold4d).
478 ///
479 /// Composes [conv_transpose2d](ModuleOps::conv_transpose2d) with the same one-hot weight
480 /// [unfold4d](ModuleOps::unfold4d) uses, so backends inherit a correct (and differentiable)
481 /// implementation for free and may override it with a custom one.
482 ///
483 /// # Shapes
484 ///
485 /// x: `[batch_size, channels * kernel_size_0 * kernel_size_1, num_blocks]`,
486 /// output: `[batch_size, channels, output_size_0, output_size_1]`
487 fn fold4d(
488 x: FloatTensor<B>,
489 output_size: [usize; 2],
490 kernel_size: [usize; 2],
491 options: UnfoldOptions,
492 ) -> FloatTensor<B> {
493 let [batch_size, channels_col, num_blocks] = x.shape().dims();
494 let [kernel_height, kernel_width] = kernel_size;
495 let [output_height, output_width] = output_size;
496 let [stride_height, stride_width] = options.stride;
497 let [padding_height, padding_width] = options.padding;
498 let [dilation_height, dilation_width] = options.dilation;
499
500 let kernel_elems = kernel_height * kernel_width;
501 assert_eq!(
502 channels_col % kernel_elems,
503 0,
504 "fold4d: input channels ({channels_col}) must be divisible by the kernel size product ({kernel_elems})"
505 );
506 let channels = channels_col / kernel_elems;
507
508 // Number of sliding blocks along each spatial dimension (the unfold output grid).
509 let blocks_height =
510 (output_height + 2 * padding_height - dilation_height * (kernel_height - 1) - 1)
511 / stride_height
512 + 1;
513 let blocks_width =
514 (output_width + 2 * padding_width - dilation_width * (kernel_width - 1) - 1)
515 / stride_width
516 + 1;
517 assert_eq!(
518 num_blocks,
519 blocks_height * blocks_width,
520 "fold4d: number of blocks ({num_blocks}) does not match the expected grid ({blocks_height} x {blocks_width}) for the given output size and options"
521 );
522
523 // The fold weight is identical to the one `unfold4d` builds for its `conv2d` — fold is its adjoint.
524 let weight = create_unfolding_weight::<B>(channels, kernel_size, &x.device(), x.dtype());
525
526 // Reshape the columns into the spatial grid of blocks, then scatter-add them back.
527 let x = B::float_reshape(
528 x,
529 Shape::new([batch_size, channels_col, blocks_height, blocks_width]),
530 );
531
532 // `padding_out` recovers the exact requested output size (always `< stride`).
533 let padding_out = [
534 (output_height + 2 * padding_height - dilation_height * (kernel_height - 1) - 1)
535 % stride_height,
536 (output_width + 2 * padding_width - dilation_width * (kernel_width - 1) - 1)
537 % stride_width,
538 ];
539
540 B::conv_transpose2d(
541 x,
542 weight,
543 None,
544 ConvTransposeOptions::new(
545 options.stride,
546 options.padding,
547 padding_out,
548 options.dilation,
549 1,
550 ),
551 )
552 }
553
554 /// One dimensional avg pooling.
555 ///
556 /// # Shapes
557 ///
558 /// x: [batch_size, channels, length],
559 fn avg_pool1d(
560 x: FloatTensor<B>,
561 kernel_size: usize,
562 stride: usize,
563 padding: usize,
564 count_include_pad: bool,
565 ceil_mode: bool,
566 ) -> FloatTensor<B> {
567 pool::avg_pool1d_from_2d::<B>(
568 x,
569 kernel_size,
570 stride,
571 padding,
572 count_include_pad,
573 ceil_mode,
574 )
575 }
576 /// Backward pass for the [avg pooling 1d](ModuleOps::avg_pool1d) operation.
577 fn avg_pool1d_backward(
578 x: FloatTensor<B>,
579 grad: FloatTensor<B>,
580 kernel_size: usize,
581 stride: usize,
582 padding: usize,
583 count_include_pad: bool,
584 ceil_mode: bool,
585 ) -> FloatTensor<B> {
586 pool::avg_pool1d_backward_from_2d::<B>(
587 x,
588 grad,
589 kernel_size,
590 stride,
591 padding,
592 count_include_pad,
593 ceil_mode,
594 )
595 }
596 /// Two dimensional avg pooling.
597 ///
598 /// # Shapes
599 ///
600 /// x: [batch_size, channels, height, width],
601 fn avg_pool2d(
602 x: FloatTensor<B>,
603 kernel_size: [usize; 2],
604 stride: [usize; 2],
605 padding: [usize; 2],
606 count_include_pad: bool,
607 ceil_mode: bool,
608 ) -> FloatTensor<B>;
609 /// Backward pass for the [avg pooling 2d](ModuleOps::avg_pool2d) operation.
610 fn avg_pool2d_backward(
611 x: FloatTensor<B>,
612 grad: FloatTensor<B>,
613 kernel_size: [usize; 2],
614 stride: [usize; 2],
615 padding: [usize; 2],
616 count_include_pad: bool,
617 ceil_mode: bool,
618 ) -> FloatTensor<B>;
619 /// Two dimensional adaptive avg pooling.
620 ///
621 /// # Shapes
622 ///
623 /// x: [batch_size, channels, height, width],
624 fn adaptive_avg_pool2d(x: FloatTensor<B>, output_size: [usize; 2]) -> FloatTensor<B>;
625 /// Backward pass for the [adaptive avg pooling 2d](ModuleOps::adaptive_avg_pool2d) operation.
626 fn adaptive_avg_pool2d_backward(x: FloatTensor<B>, grad: FloatTensor<B>) -> FloatTensor<B>;
627 /// Three dimensional adaptive avg pooling.
628 ///
629 /// # Shapes
630 ///
631 /// x: [batch_size, channels, depth, height, width],
632 fn adaptive_avg_pool3d(x: FloatTensor<B>, output_size: [usize; 3]) -> FloatTensor<B>;
633 /// Backward pass for the [adaptive avg pooling 3d](ModuleOps::adaptive_avg_pool3d) operation.
634 fn adaptive_avg_pool3d_backward(x: FloatTensor<B>, grad: FloatTensor<B>) -> FloatTensor<B>;
635 /// One dimensional adaptive avg pooling.
636 ///
637 /// # Shapes
638 ///
639 /// x: [batch_size, channels, length],
640 fn adaptive_avg_pool1d(x: FloatTensor<B>, output_size: usize) -> FloatTensor<B> {
641 pool::adaptive_avg_pool1d_from_2d::<B>(x, output_size)
642 }
643 /// Backward pass for the [adaptive avg pooling 1d](ModuleOps::adaptive_avg_pool1d) operation.
644 fn adaptive_avg_pool1d_backward(x: FloatTensor<B>, grad: FloatTensor<B>) -> FloatTensor<B> {
645 pool::adaptive_avg_pool1d_backward_from_2d::<B>(x, grad)
646 }
647 /// One dimensional max pooling.
648 ///
649 /// # Shapes
650 ///
651 /// x: [batch_size, channels, length],
652 fn max_pool1d(
653 x: FloatTensor<B>,
654 kernel_size: usize,
655 stride: usize,
656 padding: usize,
657 dilation: usize,
658 ceil_mode: bool,
659 ) -> FloatTensor<B> {
660 pool::max_pool1d_from_2d::<B>(x, kernel_size, stride, padding, dilation, ceil_mode)
661 }
662
663 /// One dimensional max pooling with indices.
664 ///
665 /// # Shapes
666 ///
667 /// x: [batch_size, channels, height, width],
668 fn max_pool1d_with_indices(
669 x: FloatTensor<B>,
670 kernel_size: usize,
671 stride: usize,
672 padding: usize,
673 dilation: usize,
674 ceil_mode: bool,
675 indices_dtype: IntDType,
676 ) -> MaxPool1dWithIndices<B> {
677 pool::max_pool1d_with_indices_from_2d::<B>(
678 x,
679 kernel_size,
680 stride,
681 padding,
682 dilation,
683 ceil_mode,
684 indices_dtype,
685 )
686 }
687 /// Backward pass for the [max pooling 1d](ModuleOps::max_pool1d_with_indices) operation.
688 #[allow(clippy::too_many_arguments)]
689 fn max_pool1d_with_indices_backward(
690 x: FloatTensor<B>,
691 kernel_size: usize,
692 stride: usize,
693 padding: usize,
694 dilation: usize,
695 ceil_mode: bool,
696 output_grad: FloatTensor<B>,
697 indices: IntTensor<B>,
698 ) -> MaxPool1dBackward<B> {
699 pool::max_pool1d_with_indices_backward_from_2d::<B>(
700 x,
701 kernel_size,
702 stride,
703 padding,
704 dilation,
705 ceil_mode,
706 output_grad,
707 indices,
708 )
709 }
710
711 /// Two dimensional max pooling.
712 ///
713 /// # Shapes
714 ///
715 /// x: [batch_size, channels, height, width],
716 fn max_pool2d(
717 x: FloatTensor<B>,
718 kernel_size: [usize; 2],
719 stride: [usize; 2],
720 padding: [usize; 2],
721 dilation: [usize; 2],
722 ceil_mode: bool,
723 ) -> FloatTensor<B>;
724
725 /// Two dimensional max pooling with indices.
726 ///
727 /// # Shapes
728 ///
729 /// x: [batch_size, channels, height, width],
730 fn max_pool2d_with_indices(
731 x: FloatTensor<B>,
732 kernel_size: [usize; 2],
733 stride: [usize; 2],
734 padding: [usize; 2],
735 dilation: [usize; 2],
736 ceil_mode: bool,
737 indices_dtype: IntDType,
738 ) -> MaxPool2dWithIndices<B>;
739 /// Backward pass for the [max pooling 2d](ModuleOps::max_pool2d_with_indices) operation.
740 #[allow(clippy::too_many_arguments)]
741 fn max_pool2d_with_indices_backward(
742 x: FloatTensor<B>,
743 kernel_size: [usize; 2],
744 stride: [usize; 2],
745 padding: [usize; 2],
746 dilation: [usize; 2],
747 ceil_mode: bool,
748 output_grad: FloatTensor<B>,
749 indices: IntTensor<B>,
750 ) -> MaxPool2dBackward<B>;
751
752 /// Down/up samples the input.
753 ///
754 /// # Shapes
755 ///
756 /// x: `[batch_size, channels, height, width]`,
757 fn interpolate(
758 x: FloatTensor<B>,
759 output_size: [usize; 2],
760 options: InterpolateOptions,
761 ) -> FloatTensor<B>;
762
763 /// Backward pass for the [interpolate](ModuleOps::interpolate) operation.
764 fn interpolate_backward(
765 x: FloatTensor<B>,
766 grad: FloatTensor<B>,
767 output_size: [usize; 2],
768 options: InterpolateOptions,
769 ) -> FloatTensor<B>;
770
771 /// Computes scaled dot-product attention: softmax(QKᵗ * scale) · V,
772 /// where scale defaults to 1/sqrt(head_dim). Optionally applies masking,
773 /// additive bias, causal masking, and softcap to the attention scores.
774 ///
775 /// # Arguments
776 /// - `query`: Query tensor of shape `[batch_size, num_heads, seq_len_q, head_dim]`
777 /// - `key`: Key tensor of shape `[batch_size, num_heads, seq_len_k, head_dim]`
778 /// - `value`: Value tensor of shape `[batch_size, num_heads, seq_len_k, val_dim]`
779 /// - `mask`: Optional boolean mask of shape `[batch_size, num_heads, seq_len_q, seq_len_k]`,
780 /// where `true` indicates positions to mask (i.e. set to -inf before softmax).
781 /// - `attn_bias`: Optional float tensor of shape `[batch_size, num_heads, seq_len_q, seq_len_k]`
782 /// added to the attention scores before softmax (e.g. ALiBi, relative position biases).
783 /// - `options`: Additional attention options (custom scale, softcap, causal masking).
784 ///
785 /// # Returns
786 /// A tensor of shape `[batch_size, num_heads, seq_len_q, val_dim]`
787 /// representing the attended context per head.
788 ///
789 /// # Note
790 /// This implementation does not support dropout and is intended for inference or
791 /// use cases where dropout is not needed.
792 fn attention(
793 query: FloatTensor<B>,
794 key: FloatTensor<B>,
795 value: FloatTensor<B>,
796 mask: Option<BoolTensor<B>>,
797 attn_bias: Option<FloatTensor<B>>,
798 options: AttentionModuleOptions,
799 ) -> FloatTensor<B>;
800
801 /// Applies Layer Normalization over the last dimension of the input tensor.
802 ///
803 /// Computes `(x - mean) / sqrt(var + epsilon) * gamma + beta`, where `mean` and
804 /// (biased) `var` are reduced over the last axis.
805 ///
806 /// # Arguments
807 ///
808 /// * `tensor` - Input tensor of shape `[..., d_model]`.
809 /// * `gamma` - Scale tensor of shape `[d_model]`.
810 /// * `beta` - Optional bias tensor of shape `[d_model]`.
811 /// * `epsilon` - Numerical stability term added to the variance before the square root.
812 ///
813 /// # Returns
814 ///
815 /// A tensor with the same shape as `tensor`.
816 fn layer_norm(
817 tensor: FloatTensor<B>,
818 gamma: FloatTensor<B>,
819 beta: Option<FloatTensor<B>>,
820 epsilon: f64,
821 ) -> FloatTensor<B> {
822 let shape = tensor.shape();
823 let rank = shape.num_dims();
824 let last_dim = rank - 1;
825 let d_model = shape[last_dim];
826
827 let mean = B::float_mean_dim(tensor.clone(), last_dim);
828 let centered = B::float_sub(tensor, mean);
829 let var = B::float_mean_dim(B::float_mul(centered.clone(), centered.clone()), last_dim);
830 let denom = B::float_sqrt(B::float_add_scalar(var, epsilon.into()));
831 let normalized = B::float_div(centered, denom);
832
833 let broadcast_dims: alloc::vec::Vec<usize> = (0..rank)
834 .map(|i| if i == last_dim { d_model } else { 1 })
835 .collect();
836 let gamma_b = B::float_reshape(gamma, Shape::from(broadcast_dims.clone()));
837 let scaled = B::float_mul(normalized, gamma_b);
838
839 match beta {
840 Some(beta) => {
841 let beta_b = B::float_reshape(beta, Shape::from(broadcast_dims));
842 B::float_add(scaled, beta_b)
843 }
844 None => scaled,
845 }
846 }
847
848 /// Computes the Connectionist Temporal Classification (CTC) loss.
849 ///
850 /// Sums over all valid alignments between the input and target sequences
851 /// using the forward (alpha) algorithm.
852 ///
853 /// # Arguments
854 ///
855 /// * `log_probs` - Log-probabilities of shape `[T, N, C]`
856 /// * `targets` - Target label indices of shape `[N, S]`
857 /// * `input_lengths` - Actual input sequence lengths per batch element `[N]`
858 /// * `target_lengths` - Actual target lengths per batch element `[N]`
859 /// * `blank` - Index of the blank label
860 ///
861 /// # Returns
862 ///
863 /// Per-sample loss of shape `[N]`
864 fn ctc_loss(
865 log_probs: FloatTensor<B>,
866 targets: IntTensor<B>,
867 input_lengths: IntTensor<B>,
868 target_lengths: IntTensor<B>,
869 blank: usize,
870 ) -> FloatTensor<B> {
871 ctc::ctc_loss_default::<B>(log_probs, targets, input_lengths, target_lengths, blank)
872 }
873
874 /// Returns `true` if this backend implements [ctc_loss_backward](ModuleOps::ctc_loss_backward)
875 /// natively.
876 ///
877 /// Autodiff queries this flag to decide between two paths:
878 /// - `true`: use the backend's [ctc_loss](ModuleOps::ctc_loss) and
879 /// [ctc_loss_backward](ModuleOps::ctc_loss_backward) directly.
880 /// - `false`: call [ctc::ctc_loss_default] for the forward pass; autodiff
881 /// then differentiates through the decomposed tensor ops.
882 ///
883 /// Backends that override `ctc_loss_backward` must also override this to
884 /// return `true`.
885 fn has_ctc_loss_backward() -> bool {
886 false
887 }
888
889 /// Backward pass for [ctc_loss](ModuleOps::ctc_loss): gradient w.r.t. `log_probs`.
890 ///
891 /// Only called when [has_ctc_loss_backward](ModuleOps::has_ctc_loss_backward)
892 /// returns `true`. Backends without a native implementation should leave
893 /// both methods at their defaults; the gradient is computed automatically by
894 /// autodiff against the decomposed [ctc::ctc_loss_default] forward.
895 ///
896 /// # Arguments
897 ///
898 /// * `log_probs` - Log-probabilities of shape `[T, N, C]`
899 /// * `targets` - Target label indices of shape `[N, S]`
900 /// * `input_lengths` - Actual input sequence lengths per batch element `[N]`
901 /// * `target_lengths` - Actual target lengths per batch element `[N]`
902 /// * `grad_loss` - Upstream gradient w.r.t. the per-sample loss `[N]`
903 /// * `blank` - Index of the blank label
904 ///
905 /// # Returns
906 ///
907 /// Gradient w.r.t. `log_probs` of shape `[T, N, C]`
908 fn ctc_loss_backward(
909 _log_probs: FloatTensor<B>,
910 _targets: IntTensor<B>,
911 _input_lengths: IntTensor<B>,
912 _target_lengths: IntTensor<B>,
913 _grad_loss: FloatTensor<B>,
914 _blank: usize,
915 ) -> FloatTensor<B> {
916 unreachable!(
917 "ctc_loss_backward called on a backend whose has_ctc_loss_backward() returns false"
918 )
919 }
920
921 /// Real-valued FFT with optional size parameter.
922 ///
923 /// When `n` is `None`, the signal must be a power of two along `dim`, and the output has
924 /// `signal_len / 2 + 1` frequency bins.
925 ///
926 /// When `n` is `Some(size)`, `size` must also be a power of two. The signal is truncated
927 /// or zero-padded to `size` and the output has `size / 2 + 1` frequency bins. Non-power-
928 /// of-two sizes are currently rejected at the public API boundary; true arbitrary-`n` DFT
929 /// support (Bluestein's algorithm) is tracked as a follow-up.
930 ///
931 /// Returns two tensors: the real part and the imaginary part.
932 fn rfft(
933 signal: FloatTensor<B>,
934 dim: usize,
935 n: Option<usize>,
936 ) -> (FloatTensor<B>, FloatTensor<B>);
937
938 /// Inverse real-valued FFT with optional output size.
939 ///
940 /// When `n` is `None`, the reconstructed signal length `2 * (spectrum_size - 1)` must be
941 /// a power of two.
942 ///
943 /// When `n` is `Some(size)`, `size` must also be a power of two. Output has exactly
944 /// `size` samples.
945 fn irfft(
946 spectrum_re: FloatTensor<B>,
947 spectrum_im: FloatTensor<B>,
948 dim: usize,
949 n: Option<usize>,
950 ) -> FloatTensor<B>;
951}