burn_flex/ops/
module.rs

1//! Module operations for the Flex backend.
2//!
3//! These operations power neural network modules like convolutions and pooling.
4
5use crate::ops::{conv, conv_transpose, deform_conv, interpolate, pool};
6use crate::{Flex, FlexTensor, Layout};
7use burn_backend::{
8    DType, Element, TensorMetadata,
9    ops::{
10        AttentionModuleOptions, ConvOptions, ConvTransposeOptions, DeformConv2dBackward,
11        DeformConvOptions, FloatTensorOps, IntTensorOps, InterpolateMode, InterpolateOptions,
12        MaxPool2dBackward, MaxPool2dWithIndices, ModuleOps,
13    },
14    tensor::{BoolTensor, FloatTensor, IntTensor},
15};
16use burn_std::{Bytes, IntDType, Shape};
17use bytemuck::Pod;
18
19/// Cast a tensor from half-precision type E to f32.
20pub(crate) fn cast_to_f32<E: Element + Pod + Copy>(
21    tensor: FlexTensor,
22    to_f32: fn(E) -> f32,
23) -> FlexTensor {
24    let tensor = tensor.to_contiguous();
25    let shape = tensor.layout().shape().clone();
26    let data: &[E] = tensor.storage();
27    let f32_data: alloc::vec::Vec<f32> = data.iter().map(|&v| to_f32(v)).collect();
28    let bytes = Bytes::from_elems(f32_data);
29    FlexTensor::new(bytes, Layout::contiguous(shape), DType::F32)
30}
31
32/// Cast a tensor from f32 back to half-precision type E.
33pub(crate) fn cast_from_f32<E: Element + Pod + Copy>(
34    tensor: FlexTensor,
35    from_f32: fn(f32) -> E,
36) -> FlexTensor {
37    let tensor = tensor.to_contiguous();
38    let shape = tensor.layout().shape().clone();
39    let data: &[f32] = tensor.storage();
40    let half_data: alloc::vec::Vec<E> = data.iter().map(|&v| from_f32(v)).collect();
41    let bytes = Bytes::from_elems(half_data);
42    FlexTensor::new(bytes, Layout::contiguous(shape), E::dtype())
43}
44
45impl ModuleOps<Flex> for Flex {
46    fn conv1d(
47        x: FloatTensor<Flex>,
48        weight: FloatTensor<Flex>,
49        bias: Option<FloatTensor<Flex>>,
50        options: ConvOptions<1>,
51    ) -> FloatTensor<Flex> {
52        match x.dtype() {
53            DType::F32 => conv::conv1d_f32(x, weight, bias, &options),
54            DType::F64 => conv::conv1d_f64(x, weight, bias, &options),
55            DType::F16 => conv::conv1d_f16(x, weight, bias, &options),
56            DType::BF16 => conv::conv1d_bf16(x, weight, bias, &options),
57            dtype => panic!("conv1d: unsupported dtype {:?}", dtype),
58        }
59    }
60
61    fn conv2d(
62        x: FloatTensor<Flex>,
63        weight: FloatTensor<Flex>,
64        bias: Option<FloatTensor<Flex>>,
65        options: ConvOptions<2>,
66    ) -> FloatTensor<Flex> {
67        match x.dtype() {
68            DType::F32 => conv::conv2d_f32(x, weight, bias, &options),
69            DType::F64 => conv::conv2d_f64(x, weight, bias, &options),
70            DType::F16 => conv::conv2d_f16(x, weight, bias, &options),
71            DType::BF16 => conv::conv2d_bf16(x, weight, bias, &options),
72            dtype => panic!("conv2d: unsupported dtype {:?}", dtype),
73        }
74    }
75
76    fn deform_conv2d(
77        x: FloatTensor<Flex>,
78        offset: FloatTensor<Flex>,
79        weight: FloatTensor<Flex>,
80        mask: Option<FloatTensor<Flex>>,
81        bias: Option<FloatTensor<Flex>>,
82        options: DeformConvOptions<2>,
83    ) -> FloatTensor<Flex> {
84        match x.dtype() {
85            DType::F32 => deform_conv::deform_conv2d_f32(
86                x,
87                offset,
88                weight,
89                mask,
90                bias,
91                options.stride,
92                options.padding,
93                options.dilation,
94                options.weight_groups,
95                options.offset_groups,
96            ),
97            DType::F64 => deform_conv::deform_conv2d_f64(
98                x,
99                offset,
100                weight,
101                mask,
102                bias,
103                options.stride,
104                options.padding,
105                options.dilation,
106                options.weight_groups,
107                options.offset_groups,
108            ),
109            DType::F16 => {
110                use burn_std::f16;
111                let result = deform_conv::deform_conv2d_f32(
112                    cast_to_f32(x, f16::to_f32),
113                    cast_to_f32(offset, f16::to_f32),
114                    cast_to_f32(weight, f16::to_f32),
115                    mask.map(|m| cast_to_f32(m, f16::to_f32)),
116                    bias.map(|b| cast_to_f32(b, f16::to_f32)),
117                    options.stride,
118                    options.padding,
119                    options.dilation,
120                    options.weight_groups,
121                    options.offset_groups,
122                );
123                cast_from_f32(result, f16::from_f32)
124            }
125            DType::BF16 => {
126                use burn_std::bf16;
127                let result = deform_conv::deform_conv2d_f32(
128                    cast_to_f32(x, bf16::to_f32),
129                    cast_to_f32(offset, bf16::to_f32),
130                    cast_to_f32(weight, bf16::to_f32),
131                    mask.map(|m| cast_to_f32(m, bf16::to_f32)),
132                    bias.map(|b| cast_to_f32(b, bf16::to_f32)),
133                    options.stride,
134                    options.padding,
135                    options.dilation,
136                    options.weight_groups,
137                    options.offset_groups,
138                );
139                cast_from_f32(result, bf16::from_f32)
140            }
141            dtype => panic!("deform_conv2d: unsupported dtype {:?}", dtype),
142        }
143    }
144
145    fn deform_conv2d_backward(
146        x: FloatTensor<Flex>,
147        offset: FloatTensor<Flex>,
148        weight: FloatTensor<Flex>,
149        mask: Option<FloatTensor<Flex>>,
150        bias: Option<FloatTensor<Flex>>,
151        output_grad: FloatTensor<Flex>,
152        options: DeformConvOptions<2>,
153    ) -> DeformConv2dBackward<Flex> {
154        let (x_grad, offset_grad, weight_grad, mask_grad, bias_grad) = match x.dtype() {
155            DType::F32 => deform_conv::deform_conv2d_backward_f32(
156                x,
157                offset,
158                weight,
159                mask,
160                bias,
161                output_grad,
162                options.stride,
163                options.padding,
164                options.dilation,
165                options.weight_groups,
166                options.offset_groups,
167            ),
168            DType::F16 => {
169                use burn_std::f16;
170                let (xg, og, wg, mg, bg) = deform_conv::deform_conv2d_backward_f32(
171                    cast_to_f32(x, f16::to_f32),
172                    cast_to_f32(offset, f16::to_f32),
173                    cast_to_f32(weight, f16::to_f32),
174                    mask.map(|m| cast_to_f32(m, f16::to_f32)),
175                    bias.map(|b| cast_to_f32(b, f16::to_f32)),
176                    cast_to_f32(output_grad, f16::to_f32),
177                    options.stride,
178                    options.padding,
179                    options.dilation,
180                    options.weight_groups,
181                    options.offset_groups,
182                );
183                (
184                    cast_from_f32(xg, f16::from_f32),
185                    cast_from_f32(og, f16::from_f32),
186                    cast_from_f32(wg, f16::from_f32),
187                    mg.map(|m| cast_from_f32(m, f16::from_f32)),
188                    bg.map(|b| cast_from_f32(b, f16::from_f32)),
189                )
190            }
191            DType::BF16 => {
192                use burn_std::bf16;
193                let (xg, og, wg, mg, bg) = deform_conv::deform_conv2d_backward_f32(
194                    cast_to_f32(x, bf16::to_f32),
195                    cast_to_f32(offset, bf16::to_f32),
196                    cast_to_f32(weight, bf16::to_f32),
197                    mask.map(|m| cast_to_f32(m, bf16::to_f32)),
198                    bias.map(|b| cast_to_f32(b, bf16::to_f32)),
199                    cast_to_f32(output_grad, bf16::to_f32),
200                    options.stride,
201                    options.padding,
202                    options.dilation,
203                    options.weight_groups,
204                    options.offset_groups,
205                );
206                (
207                    cast_from_f32(xg, bf16::from_f32),
208                    cast_from_f32(og, bf16::from_f32),
209                    cast_from_f32(wg, bf16::from_f32),
210                    mg.map(|m| cast_from_f32(m, bf16::from_f32)),
211                    bg.map(|b| cast_from_f32(b, bf16::from_f32)),
212                )
213            }
214            // f64 backward computed via f32: precision loss for large/small values.
215            // A native f64 implementation would require duplicating ~400 lines of
216            // deform_conv2d_backward. f64 deform_conv is rare in practice.
217            DType::F64 => {
218                let to = |v: f64| v as f32;
219                let from = |v: f32| v as f64;
220                let (xg, og, wg, mg, bg) = deform_conv::deform_conv2d_backward_f32(
221                    cast_to_f32(x, to),
222                    cast_to_f32(offset, to),
223                    cast_to_f32(weight, to),
224                    mask.map(|m| cast_to_f32(m, to)),
225                    bias.map(|b| cast_to_f32(b, to)),
226                    cast_to_f32(output_grad, to),
227                    options.stride,
228                    options.padding,
229                    options.dilation,
230                    options.weight_groups,
231                    options.offset_groups,
232                );
233                (
234                    cast_from_f32(xg, from),
235                    cast_from_f32(og, from),
236                    cast_from_f32(wg, from),
237                    mg.map(|m| cast_from_f32(m, from)),
238                    bg.map(|b| cast_from_f32(b, from)),
239                )
240            }
241            dtype => panic!("deform_conv2d_backward: unsupported dtype {:?}", dtype),
242        };
243        DeformConv2dBackward::new(x_grad, offset_grad, weight_grad, mask_grad, bias_grad)
244    }
245
246    fn conv3d(
247        x: FloatTensor<Flex>,
248        weight: FloatTensor<Flex>,
249        bias: Option<FloatTensor<Flex>>,
250        options: ConvOptions<3>,
251    ) -> FloatTensor<Flex> {
252        match x.dtype() {
253            DType::F32 => conv::conv3d_f32(x, weight, bias, &options),
254            DType::F64 => conv::conv3d_f64(x, weight, bias, &options),
255            DType::F16 => conv::conv3d_f16(x, weight, bias, &options),
256            DType::BF16 => conv::conv3d_bf16(x, weight, bias, &options),
257            dtype => panic!("conv3d: unsupported dtype {:?}", dtype),
258        }
259    }
260
261    fn conv_transpose1d(
262        x: FloatTensor<Flex>,
263        weight: FloatTensor<Flex>,
264        bias: Option<FloatTensor<Flex>>,
265        options: ConvTransposeOptions<1>,
266    ) -> FloatTensor<Flex> {
267        match x.dtype() {
268            DType::F32 => conv_transpose::conv_transpose1d_f32(x, weight, bias, &options),
269            DType::F64 => conv_transpose::conv_transpose1d_f64(x, weight, bias, &options),
270            DType::F16 => conv_transpose::conv_transpose1d_f16(x, weight, bias, &options),
271            DType::BF16 => conv_transpose::conv_transpose1d_bf16(x, weight, bias, &options),
272            dtype => panic!("conv_transpose1d: unsupported dtype {:?}", dtype),
273        }
274    }
275
276    fn conv_transpose2d(
277        x: FloatTensor<Flex>,
278        weight: FloatTensor<Flex>,
279        bias: Option<FloatTensor<Flex>>,
280        options: ConvTransposeOptions<2>,
281    ) -> FloatTensor<Flex> {
282        match x.dtype() {
283            DType::F32 => conv_transpose::conv_transpose2d_f32(x, weight, bias, &options),
284            DType::F64 => conv_transpose::conv_transpose2d_f64(x, weight, bias, &options),
285            DType::F16 => conv_transpose::conv_transpose2d_f16(x, weight, bias, &options),
286            DType::BF16 => conv_transpose::conv_transpose2d_bf16(x, weight, bias, &options),
287            dtype => panic!("conv_transpose2d: unsupported dtype {:?}", dtype),
288        }
289    }
290
291    fn conv_transpose3d(
292        x: FloatTensor<Flex>,
293        weight: FloatTensor<Flex>,
294        bias: Option<FloatTensor<Flex>>,
295        options: ConvTransposeOptions<3>,
296    ) -> FloatTensor<Flex> {
297        match x.dtype() {
298            DType::F32 => conv_transpose::conv_transpose3d_f32(x, weight, bias, &options),
299            DType::F64 => conv_transpose::conv_transpose3d_f64(x, weight, bias, &options),
300            DType::F16 => conv_transpose::conv_transpose3d_f16(x, weight, bias, &options),
301            DType::BF16 => conv_transpose::conv_transpose3d_bf16(x, weight, bias, &options),
302            dtype => panic!("conv_transpose3d: unsupported dtype {:?}", dtype),
303        }
304    }
305
306    fn avg_pool2d(
307        x: FloatTensor<Flex>,
308        kernel_size: [usize; 2],
309        stride: [usize; 2],
310        padding: [usize; 2],
311        count_include_pad: bool,
312        ceil_mode: bool,
313    ) -> FloatTensor<Flex> {
314        match x.dtype() {
315            DType::F32 => pool::avg_pool2d_f32(
316                x,
317                kernel_size,
318                stride,
319                padding,
320                count_include_pad,
321                ceil_mode,
322            ),
323            DType::F64 => pool::avg_pool2d_f64(
324                x,
325                kernel_size,
326                stride,
327                padding,
328                count_include_pad,
329                ceil_mode,
330            ),
331            DType::F16 => pool::avg_pool2d_f16(
332                x,
333                kernel_size,
334                stride,
335                padding,
336                count_include_pad,
337                ceil_mode,
338            ),
339            DType::BF16 => pool::avg_pool2d_bf16(
340                x,
341                kernel_size,
342                stride,
343                padding,
344                count_include_pad,
345                ceil_mode,
346            ),
347            dtype => panic!("avg_pool2d: unsupported dtype {:?}", dtype),
348        }
349    }
350
351    fn avg_pool2d_backward(
352        x: FloatTensor<Flex>,
353        grad: FloatTensor<Flex>,
354        kernel_size: [usize; 2],
355        stride: [usize; 2],
356        padding: [usize; 2],
357        count_include_pad: bool,
358        _divisor_override: bool,
359    ) -> FloatTensor<Flex> {
360        match x.dtype() {
361            DType::F32 => pool::avg_pool2d_backward_f32(
362                x,
363                grad,
364                kernel_size,
365                stride,
366                padding,
367                count_include_pad,
368            ),
369            DType::F64 => pool::avg_pool2d_backward_f64(
370                x,
371                grad,
372                kernel_size,
373                stride,
374                padding,
375                count_include_pad,
376            ),
377            DType::F16 => pool::avg_pool2d_backward_f16(
378                x,
379                grad,
380                kernel_size,
381                stride,
382                padding,
383                count_include_pad,
384            ),
385            DType::BF16 => pool::avg_pool2d_backward_bf16(
386                x,
387                grad,
388                kernel_size,
389                stride,
390                padding,
391                count_include_pad,
392            ),
393            dtype => panic!("avg_pool2d_backward: unsupported dtype {:?}", dtype),
394        }
395    }
396
397    fn adaptive_avg_pool2d(x: FloatTensor<Flex>, output_size: [usize; 2]) -> FloatTensor<Flex> {
398        match x.dtype() {
399            DType::F32 => pool::adaptive_avg_pool2d_f32(x, output_size),
400            DType::F64 => pool::adaptive_avg_pool2d_f64(x, output_size),
401            DType::F16 => pool::adaptive_avg_pool2d_f16(x, output_size),
402            DType::BF16 => pool::adaptive_avg_pool2d_bf16(x, output_size),
403            dtype => panic!("adaptive_avg_pool2d: unsupported dtype {:?}", dtype),
404        }
405    }
406
407    fn adaptive_avg_pool2d_backward(
408        x: FloatTensor<Flex>,
409        grad: FloatTensor<Flex>,
410    ) -> FloatTensor<Flex> {
411        match x.dtype() {
412            DType::F32 => pool::adaptive_avg_pool2d_backward_f32(x, grad),
413            DType::F64 => pool::adaptive_avg_pool2d_backward_f64(x, grad),
414            DType::F16 => pool::adaptive_avg_pool2d_backward_f16(x, grad),
415            DType::BF16 => pool::adaptive_avg_pool2d_backward_bf16(x, grad),
416            dtype => panic!(
417                "adaptive_avg_pool2d_backward: unsupported dtype {:?}",
418                dtype
419            ),
420        }
421    }
422
423    fn adaptive_avg_pool3d(x: FloatTensor<Flex>, output_size: [usize; 3]) -> FloatTensor<Flex> {
424        match x.dtype() {
425            DType::F32 => pool::adaptive_avg_pool3d_f32(x, output_size),
426            DType::F64 => pool::adaptive_avg_pool3d_f64(x, output_size),
427            DType::F16 => pool::adaptive_avg_pool3d_f16(x, output_size),
428            DType::BF16 => pool::adaptive_avg_pool3d_bf16(x, output_size),
429            dtype => panic!("adaptive_avg_pool3d: unsupported dtype {:?}", dtype),
430        }
431    }
432
433    fn adaptive_avg_pool3d_backward(
434        x: FloatTensor<Flex>,
435        grad: FloatTensor<Flex>,
436    ) -> FloatTensor<Flex> {
437        match x.dtype() {
438            DType::F32 => pool::adaptive_avg_pool3d_backward_f32(x, grad),
439            DType::F64 => pool::adaptive_avg_pool3d_backward_f64(x, grad),
440            DType::F16 => pool::adaptive_avg_pool3d_backward_f16(x, grad),
441            DType::BF16 => pool::adaptive_avg_pool3d_backward_bf16(x, grad),
442            dtype => panic!(
443                "adaptive_avg_pool3d_backward: unsupported dtype {:?}",
444                dtype
445            ),
446        }
447    }
448
449    fn max_pool2d(
450        x: FloatTensor<Flex>,
451        kernel_size: [usize; 2],
452        stride: [usize; 2],
453        padding: [usize; 2],
454        dilation: [usize; 2],
455        ceil_mode: bool,
456    ) -> FloatTensor<Flex> {
457        match x.dtype() {
458            DType::F32 => {
459                pool::max_pool2d_f32(x, kernel_size, stride, padding, dilation, ceil_mode)
460            }
461            DType::F64 => {
462                pool::max_pool2d_f64(x, kernel_size, stride, padding, dilation, ceil_mode)
463            }
464            DType::F16 => {
465                pool::max_pool2d_f16(x, kernel_size, stride, padding, dilation, ceil_mode)
466            }
467            DType::BF16 => {
468                pool::max_pool2d_bf16(x, kernel_size, stride, padding, dilation, ceil_mode)
469            }
470            dtype => panic!("max_pool2d: unsupported dtype {:?}", dtype),
471        }
472    }
473
474    fn max_pool2d_with_indices(
475        x: FloatTensor<Flex>,
476        kernel_size: [usize; 2],
477        stride: [usize; 2],
478        padding: [usize; 2],
479        dilation: [usize; 2],
480        ceil_mode: bool,
481        indices_dtype: IntDType,
482    ) -> MaxPool2dWithIndices<Flex> {
483        let (output, mut indices) = match x.dtype() {
484            DType::F32 => pool::max_pool2d_with_indices_f32(
485                x,
486                kernel_size,
487                stride,
488                padding,
489                dilation,
490                ceil_mode,
491            ),
492            DType::F64 => pool::max_pool2d_with_indices_f64(
493                x,
494                kernel_size,
495                stride,
496                padding,
497                dilation,
498                ceil_mode,
499            ),
500            DType::F16 => pool::max_pool2d_with_indices_f16(
501                x,
502                kernel_size,
503                stride,
504                padding,
505                dilation,
506                ceil_mode,
507            ),
508            DType::BF16 => pool::max_pool2d_with_indices_bf16(
509                x,
510                kernel_size,
511                stride,
512                padding,
513                dilation,
514                ceil_mode,
515            ),
516            dtype => panic!("max_pool2d_with_indices: unsupported dtype {:?}", dtype),
517        };
518        if indices.dtype() != DType::from(indices_dtype) {
519            indices = Flex::int_cast(indices, indices_dtype);
520        }
521        MaxPool2dWithIndices::new(output, indices)
522    }
523
524    fn max_pool2d_with_indices_backward(
525        x: FloatTensor<Flex>,
526        _kernel_size: [usize; 2],
527        _stride: [usize; 2],
528        _padding: [usize; 2],
529        _dilation: [usize; 2],
530        _ceil_mode: bool,
531        output_grad: FloatTensor<Flex>,
532        indices: IntTensor<Flex>,
533    ) -> MaxPool2dBackward<Flex> {
534        let x_grad = match x.dtype() {
535            DType::F32 => pool::max_pool2d_backward_f32(x, output_grad, indices),
536            DType::F64 => pool::max_pool2d_backward_f64(x, output_grad, indices),
537            DType::F16 => pool::max_pool2d_backward_f16(x, output_grad, indices),
538            DType::BF16 => pool::max_pool2d_backward_bf16(x, output_grad, indices),
539            dtype => panic!(
540                "max_pool2d_with_indices_backward: unsupported dtype {:?}",
541                dtype
542            ),
543        };
544        MaxPool2dBackward::new(x_grad)
545    }
546
547    fn interpolate(
548        x: FloatTensor<Flex>,
549        output_size: [usize; 2],
550        options: InterpolateOptions,
551    ) -> FloatTensor<Flex> {
552        match (options.mode, x.dtype()) {
553            (InterpolateMode::Nearest, DType::F32) => {
554                interpolate::interpolate_nearest_f32(x, output_size, options.align_corners)
555            }
556            (InterpolateMode::Nearest, DType::F64) => {
557                interpolate::interpolate_nearest_f64(x, output_size, options.align_corners)
558            }
559            (InterpolateMode::Nearest, DType::F16) => {
560                interpolate::interpolate_nearest_f16(x, output_size, options.align_corners)
561            }
562            (InterpolateMode::Nearest, DType::BF16) => {
563                interpolate::interpolate_nearest_bf16(x, output_size, options.align_corners)
564            }
565            (InterpolateMode::Bilinear, DType::F32) => {
566                interpolate::interpolate_bilinear_f32(x, output_size, options.align_corners)
567            }
568            (InterpolateMode::Bilinear, DType::F64) => {
569                interpolate::interpolate_bilinear_f64(x, output_size, options.align_corners)
570            }
571            (InterpolateMode::Bilinear, DType::F16) => {
572                interpolate::interpolate_bilinear_f16(x, output_size, options.align_corners)
573            }
574            (InterpolateMode::Bilinear, DType::BF16) => {
575                interpolate::interpolate_bilinear_bf16(x, output_size, options.align_corners)
576            }
577            (InterpolateMode::Bicubic, DType::F32) => {
578                interpolate::interpolate_bicubic_f32(x, output_size, options.align_corners)
579            }
580            (InterpolateMode::Bicubic, DType::F64) => {
581                interpolate::interpolate_bicubic_f64(x, output_size, options.align_corners)
582            }
583            (InterpolateMode::Bicubic, DType::F16) => {
584                interpolate::interpolate_bicubic_f16(x, output_size, options.align_corners)
585            }
586            (InterpolateMode::Bicubic, DType::BF16) => {
587                interpolate::interpolate_bicubic_bf16(x, output_size, options.align_corners)
588            }
589            (InterpolateMode::Lanczos3, DType::F32) => {
590                interpolate::interpolate_lanczos3_f32(x, output_size, options.align_corners)
591            }
592            (InterpolateMode::Lanczos3, DType::F64) => {
593                interpolate::interpolate_lanczos3_f64(x, output_size, options.align_corners)
594            }
595            (InterpolateMode::Lanczos3, DType::F16) => {
596                interpolate::interpolate_lanczos3_f16(x, output_size, options.align_corners)
597            }
598            (InterpolateMode::Lanczos3, DType::BF16) => {
599                interpolate::interpolate_lanczos3_bf16(x, output_size, options.align_corners)
600            }
601            (mode, dtype) => panic!(
602                "interpolate: unsupported mode {:?} / dtype {:?}",
603                mode, dtype
604            ),
605        }
606    }
607
608    fn interpolate_backward(
609        x: FloatTensor<Flex>,
610        grad: FloatTensor<Flex>,
611        output_size: [usize; 2],
612        options: InterpolateOptions,
613    ) -> FloatTensor<Flex> {
614        match (options.mode, x.dtype()) {
615            (InterpolateMode::Nearest, DType::F32) => {
616                interpolate::interpolate_nearest_backward_f32(
617                    x,
618                    grad,
619                    output_size,
620                    options.align_corners,
621                )
622            }
623            (InterpolateMode::Nearest, DType::F64) => {
624                interpolate::interpolate_nearest_backward_f64(
625                    x,
626                    grad,
627                    output_size,
628                    options.align_corners,
629                )
630            }
631            (InterpolateMode::Nearest, DType::F16) => {
632                interpolate::interpolate_nearest_backward_f16(
633                    x,
634                    grad,
635                    output_size,
636                    options.align_corners,
637                )
638            }
639            (InterpolateMode::Nearest, DType::BF16) => {
640                interpolate::interpolate_nearest_backward_bf16(
641                    x,
642                    grad,
643                    output_size,
644                    options.align_corners,
645                )
646            }
647            (InterpolateMode::Bilinear, DType::F32) => {
648                interpolate::interpolate_bilinear_backward_f32(
649                    x,
650                    grad,
651                    output_size,
652                    options.align_corners,
653                )
654            }
655            (InterpolateMode::Bilinear, DType::F64) => {
656                interpolate::interpolate_bilinear_backward_f64(
657                    x,
658                    grad,
659                    output_size,
660                    options.align_corners,
661                )
662            }
663            (InterpolateMode::Bilinear, DType::F16) => {
664                interpolate::interpolate_bilinear_backward_f16(
665                    x,
666                    grad,
667                    output_size,
668                    options.align_corners,
669                )
670            }
671            (InterpolateMode::Bilinear, DType::BF16) => {
672                interpolate::interpolate_bilinear_backward_bf16(
673                    x,
674                    grad,
675                    output_size,
676                    options.align_corners,
677                )
678            }
679            (InterpolateMode::Bicubic, DType::F32) => {
680                interpolate::interpolate_bicubic_backward_f32(
681                    x,
682                    grad,
683                    output_size,
684                    options.align_corners,
685                )
686            }
687            (InterpolateMode::Bicubic, DType::F64) => {
688                interpolate::interpolate_bicubic_backward_f64(
689                    x,
690                    grad,
691                    output_size,
692                    options.align_corners,
693                )
694            }
695            (InterpolateMode::Bicubic, DType::F16) => {
696                interpolate::interpolate_bicubic_backward_f16(
697                    x,
698                    grad,
699                    output_size,
700                    options.align_corners,
701                )
702            }
703            (InterpolateMode::Bicubic, DType::BF16) => {
704                interpolate::interpolate_bicubic_backward_bf16(
705                    x,
706                    grad,
707                    output_size,
708                    options.align_corners,
709                )
710            }
711            (mode, dtype) => {
712                panic!(
713                    "interpolate_backward: unsupported mode {:?} / dtype {:?}",
714                    mode, dtype
715                )
716            }
717        }
718    }
719
720    fn attention(
721        query: FloatTensor<Flex>,
722        key: FloatTensor<Flex>,
723        value: FloatTensor<Flex>,
724        mask: Option<BoolTensor<Flex>>,
725        attn_bias: Option<FloatTensor<Flex>>,
726        options: AttentionModuleOptions,
727    ) -> FloatTensor<Flex> {
728        crate::ops::attention::attention(query, key, value, mask, attn_bias, options)
729    }
730
731    fn rfft(
732        signal: FloatTensor<Flex>,
733        dim: usize,
734        n: Option<usize>,
735    ) -> (FloatTensor<Flex>, FloatTensor<Flex>) {
736        match signal.dtype() {
737            DType::F32 => crate::ops::fft::rfft_f32(signal, dim, n),
738            DType::F64 => crate::ops::fft::rfft_f64(signal, dim, n),
739            DType::F16 => crate::ops::fft::rfft_f16(signal, dim, n),
740            DType::BF16 => crate::ops::fft::rfft_bf16(signal, dim, n),
741            dtype => panic!("rfft: unsupported dtype {:?}", dtype),
742        }
743    }
744
745    fn irfft(
746        spectrum_re: FloatTensor<Flex>,
747        spectrum_im: FloatTensor<Flex>,
748        dim: usize,
749        n: Option<usize>,
750    ) -> FloatTensor<Flex> {
751        match spectrum_re.dtype() {
752            DType::F32 => crate::ops::fft::irfft_f32(spectrum_re, spectrum_im, dim, n),
753            DType::F64 => crate::ops::fft::irfft_f64(spectrum_re, spectrum_im, dim, n),
754            DType::F16 => crate::ops::fft::irfft_f16(spectrum_re, spectrum_im, dim, n),
755            DType::BF16 => crate::ops::fft::irfft_bf16(spectrum_re, spectrum_im, dim, n),
756            dtype => panic!("irfft: unsupported dtype {:?}", dtype),
757        }
758    }
759
760    fn embedding(weights: FloatTensor<Flex>, indices: IntTensor<Flex>) -> FloatTensor<Flex> {
761        let [batch_size, seq_length] = indices.shape().dims();
762        let [_, d_model] = weights.shape().dims();
763
764        let indices = Flex::int_reshape(indices, Shape::from(alloc::vec![batch_size * seq_length]));
765        let output = Flex::float_select(weights, 0, indices);
766        Flex::float_reshape(
767            output,
768            Shape::from(alloc::vec![batch_size, seq_length, d_model]),
769        )
770    }
771
772    fn layer_norm(
773        tensor: FloatTensor<Flex>,
774        gamma: FloatTensor<Flex>,
775        beta: Option<FloatTensor<Flex>>,
776        epsilon: f64,
777    ) -> FloatTensor<Flex> {
778        crate::ops::activation::layer_norm(tensor, gamma, beta, epsilon)
779    }
780
781    fn embedding_backward(
782        weights: FloatTensor<Flex>,
783        output_grad: FloatTensor<Flex>,
784        indices: IntTensor<Flex>,
785    ) -> FloatTensor<Flex> {
786        let [batch_size, seq_length] = indices.shape().dims();
787        let [n_embeddings, d_model] = weights.shape().dims();
788        let dtype = output_grad.dtype();
789
790        let indices = Flex::int_reshape(indices, Shape::from(alloc::vec![batch_size * seq_length]));
791        let output_grad = Flex::float_reshape(
792            output_grad,
793            Shape::from(alloc::vec![batch_size * seq_length, d_model]),
794        );
795        let grad = Flex::float_zeros(
796            Shape::from(alloc::vec![n_embeddings, d_model]),
797            &Default::default(),
798            dtype.into(),
799        );
800        Flex::float_select_add(grad, 0, indices, output_grad)
801    }
802}
burn_flex/ops/module.rs

burn_flex/ops/
module.rs
