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    /// Two dimensional adaptive avg pooling.
598    ///
599    /// # Shapes
600    ///
601    /// x: [batch_size, channels, height, width],
602    fn adaptive_avg_pool2d(x: FloatTensor<B>, output_size: [usize; 2]) -> FloatTensor<B>;
603    /// Backward pass for the [adaptive avg pooling 2d](ModuleOps::adaptive_avg_pool2d) operation.
604    fn adaptive_avg_pool2d_backward(x: FloatTensor<B>, grad: FloatTensor<B>) -> FloatTensor<B>;
605    /// Three dimensional adaptive average pooling of `[batch, channels, depth, height, width]`.
606    fn adaptive_avg_pool3d(x: FloatTensor<B>, output_size: [usize; 3]) -> FloatTensor<B> {
607        pool::adaptive_avg_pool3d_from_2d::<B>(x, output_size)
608    }
609    /// Input gradients for native three dimensional adaptive average pooling.
610    fn adaptive_avg_pool3d_backward(x: FloatTensor<B>, grad: FloatTensor<B>) -> FloatTensor<B> {
611        pool::adaptive_avg_pool3d_backward_from_2d::<B>(x, grad)
612    }
613    /// One dimensional adaptive avg pooling.
614    ///
615    /// # Shapes
616    ///
617    /// x: [batch_size, channels, length],
618    fn adaptive_avg_pool1d(x: FloatTensor<B>, output_size: usize) -> FloatTensor<B> {
619        pool::adaptive_avg_pool1d_from_2d::<B>(x, output_size)
620    }
621    /// Backward pass for the [adaptive avg pooling 1d](ModuleOps::adaptive_avg_pool1d) operation.
622    fn adaptive_avg_pool1d_backward(x: FloatTensor<B>, grad: FloatTensor<B>) -> FloatTensor<B> {
623        pool::adaptive_avg_pool1d_backward_from_2d::<B>(x, grad)
624    }
625    /// One dimensional max pooling.
626    ///
627    /// # Shapes
628    ///
629    /// x: [batch_size, channels, length],
630    fn max_pool1d(
631        x: FloatTensor<B>,
632        kernel_size: usize,
633        stride: usize,
634        padding: usize,
635        dilation: usize,
636        ceil_mode: bool,
637    ) -> FloatTensor<B> {
638        pool::max_pool1d_from_2d::<B>(x, kernel_size, stride, padding, dilation, ceil_mode)
639    }
640
641    /// One dimensional max pooling with indices.
642    ///
643    /// # Shapes
644    ///
645    /// x: [batch_size, channels, height, width],
646    fn max_pool1d_with_indices(
647        x: FloatTensor<B>,
648        kernel_size: usize,
649        stride: usize,
650        padding: usize,
651        dilation: usize,
652        ceil_mode: bool,
653    ) -> MaxPool1dWithIndices<B> {
654        pool::max_pool1d_with_indices_from_2d::<B>(
655            x,
656            kernel_size,
657            stride,
658            padding,
659            dilation,
660            ceil_mode,
661        )
662    }
663    /// Backward pass for the [max pooling 1d](ModuleOps::max_pool1d_with_indices) operation.
664    #[allow(clippy::too_many_arguments)]
665    fn max_pool1d_with_indices_backward(
666        x: FloatTensor<B>,
667        kernel_size: usize,
668        stride: usize,
669        padding: usize,
670        dilation: usize,
671        ceil_mode: bool,
672        output_grad: FloatTensor<B>,
673        indices: IntTensor<B>,
674    ) -> MaxPool1dBackward<B> {
675        pool::max_pool1d_with_indices_backward_from_2d::<B>(
676            x,
677            kernel_size,
678            stride,
679            padding,
680            dilation,
681            ceil_mode,
682            output_grad,
683            indices,
684        )
685    }
686
687    /// Two dimensional max pooling.
688    ///
689    /// # Shapes
690    ///
691    /// x: [batch_size, channels, height, width],
692    fn max_pool2d(
693        x: FloatTensor<B>,
694        kernel_size: [usize; 2],
695        stride: [usize; 2],
696        padding: [usize; 2],
697        dilation: [usize; 2],
698        ceil_mode: bool,
699    ) -> FloatTensor<B>;
700
701    /// Two dimensional max pooling with indices.
702    ///
703    /// # Shapes
704    ///
705    /// x: [batch_size, channels, height, width],
706    fn max_pool2d_with_indices(
707        x: FloatTensor<B>,
708        kernel_size: [usize; 2],
709        stride: [usize; 2],
710        padding: [usize; 2],
711        dilation: [usize; 2],
712        ceil_mode: bool,
713    ) -> MaxPool2dWithIndices<B>;
714    /// Backward pass for the [max pooling 2d](ModuleOps::max_pool2d_with_indices) operation.
715    #[allow(clippy::too_many_arguments)]
716    fn max_pool2d_with_indices_backward(
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        output_grad: FloatTensor<B>,
724        indices: IntTensor<B>,
725    ) -> MaxPool2dBackward<B>;
726
727    /// Down/up samples the input.
728    ///
729    /// # Shapes
730    ///
731    /// x: `[batch_size, channels, height, width]`,
732    fn interpolate(
733        x: FloatTensor<B>,
734        output_size: [usize; 2],
735        options: InterpolateOptions,
736    ) -> FloatTensor<B>;
737
738    /// Backward pass for the [interpolate](ModuleOps::interpolate) operation.
739    fn interpolate_backward(
740        x: FloatTensor<B>,
741        grad: FloatTensor<B>,
742        output_size: [usize; 2],
743        options: InterpolateOptions,
744    ) -> FloatTensor<B>;
745
746    /// Computes scaled dot-product attention: softmax(QKᵗ * scale) · V,
747    /// where scale defaults to 1/sqrt(head_dim). Optionally applies masking,
748    /// additive bias, causal masking, and softcap to the attention scores.
749    ///
750    /// # Arguments
751    /// - `query`: Query tensor of shape `[batch_size, num_heads, seq_len_q, head_dim]`
752    /// - `key`: Key tensor of shape `[batch_size, num_heads, seq_len_k, head_dim]`
753    /// - `value`: Value tensor of shape `[batch_size, num_heads, seq_len_k, val_dim]`
754    /// - `mask`: Optional boolean mask of shape `[batch_size, num_heads, seq_len_q, seq_len_k]`,
755    ///   where `true` indicates positions to mask (i.e. set to -inf before softmax).
756    /// - `attn_bias`: Optional float tensor of shape `[batch_size, num_heads, seq_len_q, seq_len_k]`
757    ///   added to the attention scores before softmax (e.g. ALiBi, relative position biases).
758    /// - `options`: Additional attention options (custom scale, softcap, causal masking).
759    ///
760    /// # Returns
761    /// A tensor of shape `[batch_size, num_heads, seq_len_q, val_dim]`
762    /// representing the attended context per head.
763    ///
764    /// # Note
765    /// This implementation does not support dropout and is intended for inference or
766    /// use cases where dropout is not needed.
767    fn attention(
768        query: FloatTensor<B>,
769        key: FloatTensor<B>,
770        value: FloatTensor<B>,
771        mask: Option<BoolTensor<B>>,
772        attn_bias: Option<FloatTensor<B>>,
773        options: AttentionModuleOptions,
774    ) -> FloatTensor<B>;
775
776    /// Applies Layer Normalization over the last dimension of the input tensor.
777    ///
778    /// Computes `(x - mean) / sqrt(var + epsilon) * gamma + beta`, where `mean` and
779    /// (biased) `var` are reduced over the last axis.
780    ///
781    /// # Arguments
782    ///
783    /// * `tensor` - Input tensor of shape `[..., d_model]`.
784    /// * `gamma` - Scale tensor of shape `[d_model]`.
785    /// * `beta` - Optional bias tensor of shape `[d_model]`.
786    /// * `epsilon` - Numerical stability term added to the variance before the square root.
787    ///
788    /// # Returns
789    ///
790    /// A tensor with the same shape as `tensor`.
791    fn layer_norm(
792        tensor: FloatTensor<B>,
793        gamma: FloatTensor<B>,
794        beta: Option<FloatTensor<B>>,
795        epsilon: f64,
796    ) -> FloatTensor<B> {
797        Self::layer_norm_default(tensor, gamma, beta, epsilon)
798    }
799
800    /// Whether native forward statistics and complete first-order backward are available.
801    fn has_layer_norm_backward() -> bool { false }
802
803    /// Native forward with statistics retained for backward.
804    fn layer_norm_with_stats(
805        _tensor: FloatTensor<B>, _gamma: FloatTensor<B>,
806        _beta: Option<FloatTensor<B>>, _epsilon: f64,
807    ) -> LayerNormOutput<B> {
808        unimplemented!("native LayerNorm statistics unavailable")
809    }
810
811    /// Native backward using the exact statistics returned by forward.
812    fn layer_norm_backward(
813        _tensor: FloatTensor<B>, _gamma: FloatTensor<B>, _grad: FloatTensor<B>,
814        _mean: FloatTensor<B>, _rstd: FloatTensor<B>,
815    ) -> LayerNormBackward<B> {
816        unimplemented!("native LayerNorm backward unavailable")
817    }
818
819    /// Differentiable primitive composition for backends without native LayerNorm backward.
820    fn layer_norm_default(
821        tensor: FloatTensor<B>, gamma: FloatTensor<B>,
822        beta: Option<FloatTensor<B>>, epsilon: f64,
823    ) -> FloatTensor<B> {
824        let shape = tensor.shape();
825        let rank = shape.num_dims();
826        let last_dim = rank - 1;
827        let d_model = shape[last_dim];
828
829        let mean = B::float_mean_dim(tensor.clone(), last_dim);
830        let centered = B::float_sub(tensor, mean);
831        let var = B::float_mean_dim(B::float_mul(centered.clone(), centered.clone()), last_dim);
832        let denom = B::float_sqrt(B::float_add_scalar(var, epsilon.into()));
833        let normalized = B::float_div(centered, denom);
834
835        let broadcast_dims: alloc::vec::Vec<usize> = (0..rank)
836            .map(|i| if i == last_dim { d_model } else { 1 })
837            .collect();
838        let gamma_b = B::float_reshape(gamma, Shape::from(broadcast_dims.clone()));
839        let scaled = B::float_mul(normalized, gamma_b);
840
841        match beta {
842            Some(beta) => {
843                let beta_b = B::float_reshape(beta, Shape::from(broadcast_dims));
844                B::float_add(scaled, beta_b)
845            }
846            None => scaled,
847        }
848    }
849
850    /// Computes the Connectionist Temporal Classification (CTC) loss.
851    ///
852    /// Sums over all valid alignments between the input and target sequences
853    /// using the forward (alpha) algorithm.
854    ///
855    /// # Arguments
856    ///
857    /// * `log_probs` - Log-probabilities of shape `[T, N, C]`
858    /// * `targets` - Target label indices of shape `[N, S]`
859    /// * `input_lengths` - Actual input sequence lengths per batch element `[N]`
860    /// * `target_lengths` - Actual target lengths per batch element `[N]`
861    /// * `blank` - Index of the blank label
862    ///
863    /// # Returns
864    ///
865    /// Per-sample loss of shape `[N]`
866    fn ctc_loss(
867        log_probs: FloatTensor<B>,
868        targets: IntTensor<B>,
869        input_lengths: IntTensor<B>,
870        target_lengths: IntTensor<B>,
871        blank: usize,
872    ) -> FloatTensor<B> {
873        ctc::ctc_loss_default::<B>(log_probs, targets, input_lengths, target_lengths, blank)
874    }
875
876    /// Returns `true` if this backend implements [ctc_loss_backward](ModuleOps::ctc_loss_backward)
877    /// natively.
878    ///
879    /// Autodiff queries this flag to decide between two paths:
880    /// - `true`: use the backend's [ctc_loss](ModuleOps::ctc_loss) and
881    ///   [ctc_loss_backward](ModuleOps::ctc_loss_backward) directly.
882    /// - `false`: call [ctc::ctc_loss_default] for the forward pass; autodiff
883    ///   then differentiates through the decomposed tensor ops.
884    ///
885    /// Backends that override `ctc_loss_backward` must also override this to
886    /// return `true`.
887    fn has_ctc_loss_backward() -> bool {
888        false
889    }
890
891    /// Backward pass for [ctc_loss](ModuleOps::ctc_loss): gradient w.r.t. `log_probs`.
892    ///
893    /// Only called when [has_ctc_loss_backward](ModuleOps::has_ctc_loss_backward)
894    /// returns `true`. Backends without a native implementation should leave
895    /// both methods at their defaults; the gradient is computed automatically by
896    /// autodiff against the decomposed [ctc::ctc_loss_default] forward.
897    ///
898    /// # Arguments
899    ///
900    /// * `log_probs` - Log-probabilities of shape `[T, N, C]`
901    /// * `targets` - Target label indices of shape `[N, S]`
902    /// * `input_lengths` - Actual input sequence lengths per batch element `[N]`
903    /// * `target_lengths` - Actual target lengths per batch element `[N]`
904    /// * `grad_loss` - Upstream gradient w.r.t. the per-sample loss `[N]`
905    /// * `blank` - Index of the blank label
906    ///
907    /// # Returns
908    ///
909    /// Gradient w.r.t. `log_probs` of shape `[T, N, C]`
910    fn ctc_loss_backward(
911        _log_probs: FloatTensor<B>,
912        _targets: IntTensor<B>,
913        _input_lengths: IntTensor<B>,
914        _target_lengths: IntTensor<B>,
915        _grad_loss: FloatTensor<B>,
916        _blank: usize,
917    ) -> FloatTensor<B> {
918        unreachable!(
919            "ctc_loss_backward called on a backend whose has_ctc_loss_backward() returns false"
920        )
921    }
922
923    /// Real-valued FFT with optional size parameter.
924    ///
925    /// When `n` is `None`, the signal must be a power of two along `dim`, and the output has
926    /// `signal_len / 2 + 1` frequency bins.
927    ///
928    /// When `n` is `Some(size)`, `size` must also be a power of two. The signal is truncated
929    /// or zero-padded to `size` and the output has `size / 2 + 1` frequency bins. Non-power-
930    /// of-two sizes are currently rejected at the public API boundary; true arbitrary-`n` DFT
931    /// support (Bluestein's algorithm) is tracked as a follow-up.
932    ///
933    /// Returns two tensors: the real part and the imaginary part.
934    fn rfft(
935        signal: FloatTensor<B>,
936        dim: usize,
937        n: Option<usize>,
938    ) -> (FloatTensor<B>, FloatTensor<B>);
939
940    /// Inverse real-valued FFT with optional output size.
941    ///
942    /// When `n` is `None`, the reconstructed signal length `2 * (spectrum_size - 1)` must be
943    /// a power of two.
944    ///
945    /// When `n` is `Some(size)`, `size` must also be a power of two. Output has exactly
946    /// `size` samples.
947    fn irfft(
948        spectrum_re: FloatTensor<B>,
949        spectrum_im: FloatTensor<B>,
950        dim: usize,
951        n: Option<usize>,
952    ) -> FloatTensor<B>;
953}
954
955#[cfg(test)]
956mod tests {
957    use super::*;
958
959    #[test]
960    #[should_panic = "stride must be non-zero"]
961    fn conv_options_stride_zero() {
962        let _opt = ConvOptions::new([0, 1], [0, 0], [1, 1], 1);
963    }
964
965    #[test]
966    #[should_panic = "dilation must be non-zero"]
967    fn conv_options_dilation_zero() {
968        let _opt = ConvOptions::new([1, 1], [0, 0], [0, 0], 1);
969    }
970
971    #[test]
972    #[should_panic = "groups must be non-zero"]
973    fn conv_options_groups_zero() {
974        let _opt = ConvOptions::new([1, 1], [0, 0], [1, 1], 0);
975    }
976
977    #[test]
978    #[should_panic = "stride must be non-zero"]
979    fn conv_transpose_options_stride_zero() {
980        let _opt = ConvTransposeOptions::new([0, 1], [0, 0], [0, 0], [1, 1], 1);
981    }
982
983    #[test]
984    #[should_panic = "dilation must be non-zero"]
985    fn conv_transpose_options_dilation_zero() {
986        let _opt = ConvTransposeOptions::new([1, 1], [0, 0], [0, 0], [0, 0], 1);
987    }
988
989    #[test]
990    #[should_panic = "groups must be non-zero"]
991    fn conv_transpose_options_groups_zero() {
992        let _opt = ConvTransposeOptions::new([1, 1], [0, 0], [0, 0], [1, 1], 0);
993    }
994
995    #[test]
996    #[should_panic = "stride must be non-zero"]
997    fn deform_conv_options_stride_zero() {
998        let _opt = DeformConvOptions::new([0, 1], [0, 0], [1, 1], 1, 1);
999    }
1000
1001    #[test]
1002    #[should_panic = "dilation must be non-zero"]
1003    fn deform_conv_options_dilation_zero() {
1004        let _opt = DeformConvOptions::new([1, 1], [0, 0], [0, 0], 1, 1);
1005    }
1006
1007    #[test]
1008    #[should_panic = "weight groups must be non-zero"]
1009    fn deform_conv_options_weights_groups_zero() {
1010        let _opt = DeformConvOptions::new([1, 1], [0, 0], [1, 1], 0, 1);
1011    }
1012
1013    #[test]
1014    #[should_panic = "offset groups must be non-zero"]
1015    fn deform_conv_options_offset_groups_zero() {
1016        let _opt = DeformConvOptions::new([1, 1], [0, 0], [1, 1], 1, 0);
1017    }
1018
1019    #[test]
1020    #[should_panic = "stride must be non-zero"]
1021    fn unfold_options_stride_zero() {
1022        let _opt = UnfoldOptions::new([0, 1], [0, 0], [1, 1]);
1023    }
1024
1025    #[test]
1026    #[should_panic = "dilation must be non-zero"]
1027    fn unfold_options_dilation_zero() {
1028        let _opt = UnfoldOptions::new([1, 1], [0, 0], [0, 0]);
1029    }
1030}
ruda_tensor/ops/modules/base.rs

ruda_tensor/ops/modules/
base.rs
