ruda_tensor/ops/modules/
base.rs

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

ruda_tensor/ops/modules/
base.rs
