ruda_tensor/ops/modules/
base.rs

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

ruda_tensor/ops/modules/
base.rs
