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

ruda_tensor/api/
module.rs
