ruda_tensor/api/
module.rs

1use crate::api::{
2    Bool, Int, Tensor, TensorPrimitive,
3    backend::Backend,
4    check,
5    check::TensorCheck,
6    ops::{
7        AttentionModuleOptions, ConvOptions, ConvTransposeOptions, InterpolateOptions, PadMode,
8        PaddedConvOptions, UnfoldOptions,
9    },
10};
11
12use super::ops::DeformConvOptions;
13
14/// Computes the [CTC loss](crate::api::ops::ModuleOps::ctc_loss).
15///
16/// # Arguments
17///
18/// * `log_probs` - Log-probabilities of shape `[T, N, C]`
19/// * `targets` - Target label indices of shape `[N, S]`
20/// * `input_lengths` - Actual input sequence lengths per batch element `[N]`
21/// * `target_lengths` - Actual target lengths per batch element `[N]`
22/// * `blank` - Index of the blank label
23///
24/// # Returns
25///
26/// Per-sample loss of shape `[N]`
27pub fn ctc_loss<B>(
28    log_probs: Tensor<B, 3>,
29    targets: Tensor<B, 2, Int>,
30    input_lengths: Tensor<B, 1, Int>,
31    target_lengths: Tensor<B, 1, Int>,
32    blank: usize,
33) -> Tensor<B, 1>
34where
35    B: Backend,
36{
37    Tensor::new(TensorPrimitive::Float(B::ctc_loss(
38        log_probs.primitive.tensor(),
39        targets.primitive,
40        input_lengths.primitive,
41        target_lengths.primitive,
42        blank,
43    )))
44}
45
46/// Applies the [embedding module](crate::api::ops::ModuleOps::embedding).
47pub fn embedding<B>(weights: Tensor<B, 2>, indices: Tensor<B, 2, Int>) -> Tensor<B, 3>
48where
49    B: Backend,
50{
51    Tensor::new(TensorPrimitive::Float(B::embedding(
52        weights.primitive.tensor(),
53        indices.primitive,
54    )))
55}
56
57/// Applies a [1D convolution](crate::api::ops::ModuleOps::conv1d).
58///
59/// Accepts [`ConvOptions`] for symmetric padding, or [`PaddedConvOptions`] for
60/// asymmetric padding. When asymmetric padding is specified, an explicit pad
61/// operation is applied before the convolution backend op.
62pub fn conv1d<B>(
63    x: Tensor<B, 3>,
64    weight: Tensor<B, 3>,
65    bias: Option<Tensor<B, 1>>,
66    options: impl Into<PaddedConvOptions<1>>,
67) -> Tensor<B, 3>
68where
69    B: Backend,
70{
71    let padded_options = options.into();
72    check!(TensorCheck::conv(
73        "conv1d",
74        x.dims(),
75        weight.dims(),
76        padded_options.options.groups,
77    ));
78
79    if let Some(padding_end) = padded_options.padding_end {
80        let left = padded_options.options.padding[0];
81        let right = padding_end[0];
82        // For 1D (NCL format), pad the length dimension
83        let padded = x.pad((left, right, 0, 0), PadMode::Constant(0.0));
84        let zero_options = ConvOptions::new(
85            padded_options.options.stride,
86            [0],
87            padded_options.options.dilation,
88            padded_options.options.groups,
89        );
90        Tensor::new(TensorPrimitive::Float(B::conv1d(
91            padded.primitive.tensor(),
92            weight.primitive.tensor(),
93            bias.map(|b| b.primitive.tensor()),
94            zero_options,
95        )))
96    } else {
97        Tensor::new(TensorPrimitive::Float(B::conv1d(
98            x.primitive.tensor(),
99            weight.primitive.tensor(),
100            bias.map(|b| b.primitive.tensor()),
101            padded_options.options,
102        )))
103    }
104}
105
106/// Applies a [2D convolution](crate::api::ops::ModuleOps::conv2d).
107///
108/// Accepts [`ConvOptions`] for symmetric padding, or [`PaddedConvOptions`] for
109/// asymmetric padding. When asymmetric padding is specified, an explicit pad
110/// operation is applied before the convolution backend op.
111pub fn conv2d<B>(
112    x: Tensor<B, 4>,
113    weight: Tensor<B, 4>,
114    bias: Option<Tensor<B, 1>>,
115    options: impl Into<PaddedConvOptions<2>>,
116) -> Tensor<B, 4>
117where
118    B: Backend,
119{
120    let padded_options = options.into();
121    check!(TensorCheck::conv(
122        "conv2d",
123        x.dims(),
124        weight.dims(),
125        padded_options.options.groups,
126    ));
127
128    if let Some(padding_end) = padded_options.padding_end {
129        let top = padded_options.options.padding[0];
130        let left = padded_options.options.padding[1];
131        let bottom = padding_end[0];
132        let right = padding_end[1];
133        // For 2D (NCHW format), pad height and width
134        let padded = x.pad((left, right, top, bottom), PadMode::Constant(0.0));
135        let zero_options = ConvOptions::new(
136            padded_options.options.stride,
137            [0, 0],
138            padded_options.options.dilation,
139            padded_options.options.groups,
140        );
141        Tensor::new(TensorPrimitive::Float(B::conv2d(
142            padded.primitive.tensor(),
143            weight.primitive.tensor(),
144            bias.map(|b| b.primitive.tensor()),
145            zero_options,
146        )))
147    } else {
148        Tensor::new(TensorPrimitive::Float(B::conv2d(
149            x.primitive.tensor(),
150            weight.primitive.tensor(),
151            bias.map(|b| b.primitive.tensor()),
152            padded_options.options,
153        )))
154    }
155}
156
157/// Applies a [3D convolution](crate::api::ops::ModuleOps::conv3d).
158///
159/// Accepts [`ConvOptions`] for symmetric padding, or [`PaddedConvOptions`] for
160/// asymmetric padding. An explicit pad operation handles asymmetric padding
161/// before dispatching to the convolution backend.
162pub fn conv3d<B>(
163    x: Tensor<B, 5>,
164    weight: Tensor<B, 5>,
165    bias: Option<Tensor<B, 1>>,
166    options: impl Into<PaddedConvOptions<3>>,
167) -> Tensor<B, 5>
168where
169    B: Backend,
170{
171    let padded_options = options.into();
172    check!(TensorCheck::conv(
173        "conv3d",
174        x.dims(),
175        weight.dims(),
176        padded_options.options.groups,
177    ));
178
179    let mut options = padded_options.options;
180    let x = if let Some(padding_end) = padded_options.padding_end {
181        let padding: [(usize, usize); 3] =
182            core::array::from_fn(|axis| (options.padding[axis], padding_end[axis]));
183        options.padding = [0; 3];
184        x.pad(padding, PadMode::Constant(0.0))
185    } else {
186        x
187    };
188
189    Tensor::new(TensorPrimitive::Float(B::conv3d(
190        x.primitive.tensor(),
191        weight.primitive.tensor(),
192        bias.map(|b| b.primitive.tensor()),
193        options,
194    )))
195}
196
197/// Applies a [Deformable 2D convolution](crate::api::ops::ModuleOps::deform_conv2d).
198pub fn deform_conv2d<B>(
199    x: Tensor<B, 4>,
200    offset: Tensor<B, 4>,
201    weight: Tensor<B, 4>,
202    mask: Option<Tensor<B, 4>>,
203    bias: Option<Tensor<B, 1>>,
204    options: DeformConvOptions<2>,
205) -> Tensor<B, 4>
206where
207    B: Backend,
208{
209    check!(TensorCheck::conv(
210        "deform_conv2d",
211        x.dims(),
212        weight.dims(),
213        options.weight_groups,
214    ));
215    let [batch, channels, height, width] = x.dims();
216    let [_, _, kernel_height, kernel_width] = weight.dims();
217    assert!(
218        options.offset_groups > 0 && channels.is_multiple_of(options.offset_groups),
219        "deform_conv2d input channels must be divisible by non-zero offset groups"
220    );
221    let [out_height, out_width] = options.output_size(
222        [height, width], [kernel_height, kernel_width],
223    );
224    let mask_channels = options.offset_groups * kernel_height * kernel_width;
225    assert_eq!(
226        offset.dims(), [batch, 2 * mask_channels, out_height, out_width],
227        "deform_conv2d offset shape must match groups, kernel and output"
228    );
229    if let Some(mask) = mask.as_ref() {
230        assert_eq!(
231            mask.dims(), [batch, mask_channels, out_height, out_width],
232            "deform_conv2d mask shape must match groups, kernel and output"
233        );
234    }
235    Tensor::new(TensorPrimitive::Float(B::deform_conv2d(
236        x.primitive.tensor(),
237        offset.primitive.tensor(),
238        weight.primitive.tensor(),
239        mask.map(|m| m.primitive.tensor()),
240        bias.map(|b| b.primitive.tensor()),
241        options,
242    )))
243}
244
245/// Applies a [1D transposed convolution](crate::api::ops::ModuleOps::conv_transpose1d).
246pub fn conv_transpose1d<B>(
247    x: Tensor<B, 3>,
248    weight: Tensor<B, 3>,
249    bias: Option<Tensor<B, 1>>,
250    options: ConvTransposeOptions<1>,
251) -> Tensor<B, 3>
252where
253    B: Backend,
254{
255    check!(TensorCheck::conv_transpose(
256        "conv_transpose1d",
257        x.dims(),
258        weight.dims(),
259    ));
260    Tensor::new(TensorPrimitive::Float(B::conv_transpose1d(
261        x.primitive.tensor(),
262        weight.primitive.tensor(),
263        bias.map(|b| b.primitive.tensor()),
264        options,
265    )))
266}
267
268/// Applies a [2D transposed convolution](crate::api::ops::ModuleOps::conv_transpose2d).
269pub fn conv_transpose2d<B>(
270    x: Tensor<B, 4>,
271    weight: Tensor<B, 4>,
272    bias: Option<Tensor<B, 1>>,
273    options: ConvTransposeOptions<2>,
274) -> Tensor<B, 4>
275where
276    B: Backend,
277{
278    check!(TensorCheck::conv_transpose(
279        "conv_transpose2d",
280        x.dims(),
281        weight.dims(),
282    ));
283    Tensor::new(TensorPrimitive::Float(B::conv_transpose2d(
284        x.primitive.tensor(),
285        weight.primitive.tensor(),
286        bias.map(|b| b.primitive.tensor()),
287        options,
288    )))
289}
290
291/// Applies a 3D transposed convolution](crate::api::ops::ModuleOps::conv_transpose3d).
292pub fn conv_transpose3d<B>(
293    x: Tensor<B, 5>,
294    weight: Tensor<B, 5>,
295    bias: Option<Tensor<B, 1>>,
296    options: ConvTransposeOptions<3>,
297) -> Tensor<B, 5>
298where
299    B: Backend,
300{
301    check!(TensorCheck::conv_transpose(
302        "conv_transpose3d",
303        x.dims(),
304        weight.dims(),
305    ));
306    Tensor::new(TensorPrimitive::Float(B::conv_transpose3d(
307        x.primitive.tensor(),
308        weight.primitive.tensor(),
309        bias.map(|b| b.primitive.tensor()),
310        options,
311    )))
312}
313
314/// Applies a [4D to 3D unfold](crate::api::ops::ModuleOps::unfold4d).
315pub fn unfold4d<B>(x: Tensor<B, 4>, kernel_size: [usize; 2], options: UnfoldOptions) -> Tensor<B, 3>
316where
317    B: Backend,
318{
319    Tensor::new(TensorPrimitive::Float(B::unfold4d(
320        x.primitive.tensor(),
321        kernel_size,
322        options,
323    )))
324}
325
326/// Applies a [1D max pooling](crate::api::ops::ModuleOps::max_pool1d).
327pub fn max_pool1d<B>(
328    x: Tensor<B, 3>,
329    kernel_size: usize,
330    stride: usize,
331    padding: usize,
332    dilation: usize,
333    ceil_mode: bool,
334) -> Tensor<B, 3>
335where
336    B: Backend,
337{
338    Tensor::new(TensorPrimitive::Float(B::max_pool1d(
339        x.primitive.tensor(),
340        kernel_size,
341        stride,
342        padding,
343        dilation,
344        ceil_mode,
345    )))
346}
347
348/// Applies a [2D max pooling](crate::api::ops::ModuleOps::max_pool2d).
349pub fn max_pool2d<B>(
350    x: Tensor<B, 4>,
351    kernel_size: [usize; 2],
352    stride: [usize; 2],
353    padding: [usize; 2],
354    dilation: [usize; 2],
355    ceil_mode: bool,
356) -> Tensor<B, 4>
357where
358    B: Backend,
359{
360    Tensor::new(TensorPrimitive::Float(B::max_pool2d(
361        x.primitive.tensor(),
362        kernel_size,
363        stride,
364        padding,
365        dilation,
366        ceil_mode,
367    )))
368}
369
370/// Applies a [2D avg pooling](crate::api::ops::ModuleOps::avg_pool2d).
371pub fn avg_pool2d<B>(
372    x: Tensor<B, 4>,
373    kernel_size: [usize; 2],
374    stride: [usize; 2],
375    padding: [usize; 2],
376    count_include_pad: bool,
377    ceil_mode: bool,
378) -> Tensor<B, 4>
379where
380    B: Backend,
381{
382    Tensor::new(TensorPrimitive::Float(B::avg_pool2d(
383        x.primitive.tensor(),
384        kernel_size,
385        stride,
386        padding,
387        count_include_pad,
388        ceil_mode,
389    )))
390}
391
392/// Applies a [1D avg pooling](crate::api::ops::ModuleOps::avg_pool1d).
393pub fn avg_pool1d<B>(
394    x: Tensor<B, 3>,
395    kernel_size: usize,
396    stride: usize,
397    padding: usize,
398    count_include_pad: bool,
399    ceil_mode: bool,
400) -> Tensor<B, 3>
401where
402    B: Backend,
403{
404    Tensor::new(TensorPrimitive::Float(B::avg_pool1d(
405        x.primitive.tensor(),
406        kernel_size,
407        stride,
408        padding,
409        count_include_pad,
410        ceil_mode,
411    )))
412}
413
414/// Applies a [1D max pooling](crate::api::ops::ModuleOps::max_pool1d).
415pub fn max_pool1d_with_indices<B>(
416    x: Tensor<B, 3>,
417    kernel_size: usize,
418    stride: usize,
419    padding: usize,
420    dilation: usize,
421    ceil_mode: bool,
422) -> (Tensor<B, 3>, Tensor<B, 3, Int>)
423where
424    B: Backend,
425{
426    let output = B::max_pool1d_with_indices(
427        x.primitive.tensor(),
428        kernel_size,
429        stride,
430        padding,
431        dilation,
432        ceil_mode,
433    );
434
435    (
436        Tensor::new(TensorPrimitive::Float(output.output)),
437        Tensor::new(output.indices),
438    )
439}
440
441/// Applies a [2D max pooling with indices](crate::api::ops::ModuleOps::max_pool2d_with_indices).
442pub fn max_pool2d_with_indices<B>(
443    x: Tensor<B, 4>,
444    kernel_size: [usize; 2],
445    stride: [usize; 2],
446    padding: [usize; 2],
447    dilation: [usize; 2],
448    ceil_mode: bool,
449) -> (Tensor<B, 4>, Tensor<B, 4, Int>)
450where
451    B: Backend,
452{
453    let output = B::max_pool2d_with_indices(
454        x.primitive.tensor(),
455        kernel_size,
456        stride,
457        padding,
458        dilation,
459        ceil_mode,
460    );
461
462    (
463        Tensor::new(TensorPrimitive::Float(output.output)),
464        Tensor::new(output.indices),
465    )
466}
467
468/// Applies a [2D adaptive avg pooling](crate::api::ops::ModuleOps::adaptive_avg_pool2d).
469pub fn adaptive_avg_pool2d<B>(x: Tensor<B, 4>, output_size: [usize; 2]) -> Tensor<B, 4>
470where
471    B: Backend,
472{
473    Tensor::new(TensorPrimitive::Float(B::adaptive_avg_pool2d(
474        x.primitive.tensor(),
475        output_size,
476    )))
477}
478
479/// Applies a [1D adaptive avg pooling](crate::api::ops::ModuleOps::adaptive_avg_pool1d).
480pub fn adaptive_avg_pool1d<B>(x: Tensor<B, 3>, output_size: usize) -> Tensor<B, 3>
481where
482    B: Backend,
483{
484    Tensor::new(TensorPrimitive::Float(B::adaptive_avg_pool1d(
485        x.primitive.tensor(),
486        output_size,
487    )))
488}
489
490/// Applies a [2D interpolation](crate::api::ops::ModuleOps::interpolate).
491pub fn interpolate<B>(
492    x: Tensor<B, 4>,
493    output_size: [usize; 2],
494    options: InterpolateOptions,
495) -> Tensor<B, 4>
496where
497    B: Backend,
498{
499    Tensor::new(TensorPrimitive::Float(B::interpolate(
500        x.primitive.tensor(),
501        output_size,
502        options,
503    )))
504}
505
506/// Applies a linear transformation to the input tensor using the given weight and bias.
507///
508/// ```math
509/// y = x @ weight + [bias]
510/// ```
511///
512/// # Arguments:
513///
514/// - `input` is the input tensor, ``[..., d_input]``.
515/// - `weight` is the weight tensor, ``[d_input, d_output]``.
516/// - `bias` is the bias tensor (optional), ``[d_output]``.
517///
518/// # Returns:
519///
520/// The transformed tensor, ``[..., d_output]``.
521///
522/// # Compatibility
523///
524/// This function differs from PyTorch's ``torch.nn.functional.linear`` in that it does not
525/// transpose the weight matrix. In PyTorch, the weight matrix is transposed before
526/// multiplication:
527///
528/// ```math
529/// y = x @ weight^T + [bias]
530/// ```
531pub fn linear<B: Backend, const D: usize>(
532    input: Tensor<B, D>,
533    weight: Tensor<B, 2>,
534    bias: Option<Tensor<B, 1>>,
535) -> Tensor<B, D> {
536    if D == 1 {
537        // Insert and remove an extra batch dimension for the batch matmul to work.
538        let input = input.unsqueeze::<2>();
539        let output = linear(input, weight, bias);
540        return output.squeeze_dim(0);
541    }
542
543    Tensor::new(TensorPrimitive::Float(B::linear(
544        input.primitive.tensor(),
545        weight.primitive.tensor(),
546        bias.map(|b| b.primitive.tensor()),
547    )))
548}
549
550/// Computes scaled dot-product attention: softmax(QKᵗ * scale) · V,
551/// where scale defaults to 1/sqrt(head_dim) (configurable via `options.scale`).
552/// Optionally applies masking, additive bias, causal masking, and softcap.
553///
554/// # Arguments
555/// - `query`: Query tensor of shape `[batch_size, num_heads, seq_len_q, head_dim]`
556/// - `key`: Key tensor of shape `[batch_size, num_heads, seq_len_k, head_dim]`
557/// - `value`: Value tensor of shape `[batch_size, num_heads, seq_len_k, val_dim]`
558/// - `mask`: Optional boolean mask of shape `[batch_size, num_heads, seq_len_q, seq_len_k]`,
559///   where `true` indicates positions to mask (i.e. set to -inf before softmax).
560/// - `attn_bias`: Optional float tensor of shape `[batch_size, num_heads, seq_len_q, seq_len_k]`
561///   added to the attention scores before softmax (e.g. ALiBi, relative position biases).
562/// - `options`: Additional attention options (custom scale, softcap, causal masking).
563///
564/// # Returns
565/// A tensor of shape `[batch_size, num_heads, seq_len_q, val_dim]`
566/// representing the attended context per head.
567///
568/// # Note
569/// This implementation does not support dropout and is intended for inference or
570/// use cases where dropout is not needed.
571pub fn attention<B: Backend>(
572    query: Tensor<B, 4>,
573    key: Tensor<B, 4>,
574    value: Tensor<B, 4>,
575    mask: Option<Tensor<B, 4, Bool>>,
576    attn_bias: Option<Tensor<B, 4>>,
577    options: AttentionModuleOptions,
578) -> Tensor<B, 4> {
579    Tensor::new(TensorPrimitive::Float(B::attention(
580        query.primitive.tensor(),
581        key.primitive.tensor(),
582        value.primitive.tensor(),
583        mask.map(|mask| mask.primitive),
584        attn_bias.map(|bias| bias.primitive.tensor()),
585        options,
586    )))
587}
588
589/// Exports attention fallback to test backend's attention against.
590pub fn attention_fallback<B: Backend>(
591    query: Tensor<B, 4>,
592    key: Tensor<B, 4>,
593    value: Tensor<B, 4>,
594    mask: Option<Tensor<B, 4, Bool>>,
595    attn_bias: Option<Tensor<B, 4>>,
596    options: AttentionModuleOptions,
597) -> Tensor<B, 4> {
598    Tensor::new(TensorPrimitive::Float(
599        crate::api::ops::attention::attention_fallback::<B>(
600            query.primitive.tensor(),
601            key.primitive.tensor(),
602            value.primitive.tensor(),
603            mask.map(|mask| mask.primitive),
604            attn_bias.map(|bias| bias.primitive.tensor()),
605            options,
606        ),
607    ))
608}
ruda_tensor/api/module.rs

ruda_tensor/api/
module.rs
