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 max_pool2d(
424        x: FloatTensor<Flex>,
425        kernel_size: [usize; 2],
426        stride: [usize; 2],
427        padding: [usize; 2],
428        dilation: [usize; 2],
429        ceil_mode: bool,
430    ) -> FloatTensor<Flex> {
431        match x.dtype() {
432            DType::F32 => {
433                pool::max_pool2d_f32(x, kernel_size, stride, padding, dilation, ceil_mode)
434            }
435            DType::F64 => {
436                pool::max_pool2d_f64(x, kernel_size, stride, padding, dilation, ceil_mode)
437            }
438            DType::F16 => {
439                pool::max_pool2d_f16(x, kernel_size, stride, padding, dilation, ceil_mode)
440            }
441            DType::BF16 => {
442                pool::max_pool2d_bf16(x, kernel_size, stride, padding, dilation, ceil_mode)
443            }
444            dtype => panic!("max_pool2d: unsupported dtype {:?}", dtype),
445        }
446    }
447
448    fn max_pool2d_with_indices(
449        x: FloatTensor<Flex>,
450        kernel_size: [usize; 2],
451        stride: [usize; 2],
452        padding: [usize; 2],
453        dilation: [usize; 2],
454        ceil_mode: bool,
455        indices_dtype: IntDType,
456    ) -> MaxPool2dWithIndices<Flex> {
457        let (output, mut indices) = match x.dtype() {
458            DType::F32 => pool::max_pool2d_with_indices_f32(
459                x,
460                kernel_size,
461                stride,
462                padding,
463                dilation,
464                ceil_mode,
465            ),
466            DType::F64 => pool::max_pool2d_with_indices_f64(
467                x,
468                kernel_size,
469                stride,
470                padding,
471                dilation,
472                ceil_mode,
473            ),
474            DType::F16 => pool::max_pool2d_with_indices_f16(
475                x,
476                kernel_size,
477                stride,
478                padding,
479                dilation,
480                ceil_mode,
481            ),
482            DType::BF16 => pool::max_pool2d_with_indices_bf16(
483                x,
484                kernel_size,
485                stride,
486                padding,
487                dilation,
488                ceil_mode,
489            ),
490            dtype => panic!("max_pool2d_with_indices: unsupported dtype {:?}", dtype),
491        };
492        if indices.dtype() != DType::from(indices_dtype) {
493            indices = Flex::int_cast(indices, indices_dtype);
494        }
495        MaxPool2dWithIndices::new(output, indices)
496    }
497
498    fn max_pool2d_with_indices_backward(
499        x: FloatTensor<Flex>,
500        _kernel_size: [usize; 2],
501        _stride: [usize; 2],
502        _padding: [usize; 2],
503        _dilation: [usize; 2],
504        _ceil_mode: bool,
505        output_grad: FloatTensor<Flex>,
506        indices: IntTensor<Flex>,
507    ) -> MaxPool2dBackward<Flex> {
508        let x_grad = match x.dtype() {
509            DType::F32 => pool::max_pool2d_backward_f32(x, output_grad, indices),
510            DType::F64 => pool::max_pool2d_backward_f64(x, output_grad, indices),
511            DType::F16 => pool::max_pool2d_backward_f16(x, output_grad, indices),
512            DType::BF16 => pool::max_pool2d_backward_bf16(x, output_grad, indices),
513            dtype => panic!(
514                "max_pool2d_with_indices_backward: unsupported dtype {:?}",
515                dtype
516            ),
517        };
518        MaxPool2dBackward::new(x_grad)
519    }
520
521    fn interpolate(
522        x: FloatTensor<Flex>,
523        output_size: [usize; 2],
524        options: InterpolateOptions,
525    ) -> FloatTensor<Flex> {
526        match (options.mode, x.dtype()) {
527            (InterpolateMode::Nearest, DType::F32) => {
528                interpolate::interpolate_nearest_f32(x, output_size, options.align_corners)
529            }
530            (InterpolateMode::Nearest, DType::F64) => {
531                interpolate::interpolate_nearest_f64(x, output_size, options.align_corners)
532            }
533            (InterpolateMode::Nearest, DType::F16) => {
534                interpolate::interpolate_nearest_f16(x, output_size, options.align_corners)
535            }
536            (InterpolateMode::Nearest, DType::BF16) => {
537                interpolate::interpolate_nearest_bf16(x, output_size, options.align_corners)
538            }
539            (InterpolateMode::Bilinear, DType::F32) => {
540                interpolate::interpolate_bilinear_f32(x, output_size, options.align_corners)
541            }
542            (InterpolateMode::Bilinear, DType::F64) => {
543                interpolate::interpolate_bilinear_f64(x, output_size, options.align_corners)
544            }
545            (InterpolateMode::Bilinear, DType::F16) => {
546                interpolate::interpolate_bilinear_f16(x, output_size, options.align_corners)
547            }
548            (InterpolateMode::Bilinear, DType::BF16) => {
549                interpolate::interpolate_bilinear_bf16(x, output_size, options.align_corners)
550            }
551            (InterpolateMode::Bicubic, DType::F32) => {
552                interpolate::interpolate_bicubic_f32(x, output_size, options.align_corners)
553            }
554            (InterpolateMode::Bicubic, DType::F64) => {
555                interpolate::interpolate_bicubic_f64(x, output_size, options.align_corners)
556            }
557            (InterpolateMode::Bicubic, DType::F16) => {
558                interpolate::interpolate_bicubic_f16(x, output_size, options.align_corners)
559            }
560            (InterpolateMode::Bicubic, DType::BF16) => {
561                interpolate::interpolate_bicubic_bf16(x, output_size, options.align_corners)
562            }
563            (InterpolateMode::Lanczos3, DType::F32) => {
564                interpolate::interpolate_lanczos3_f32(x, output_size, options.align_corners)
565            }
566            (InterpolateMode::Lanczos3, DType::F64) => {
567                interpolate::interpolate_lanczos3_f64(x, output_size, options.align_corners)
568            }
569            (InterpolateMode::Lanczos3, DType::F16) => {
570                interpolate::interpolate_lanczos3_f16(x, output_size, options.align_corners)
571            }
572            (InterpolateMode::Lanczos3, DType::BF16) => {
573                interpolate::interpolate_lanczos3_bf16(x, output_size, options.align_corners)
574            }
575            (mode, dtype) => panic!(
576                "interpolate: unsupported mode {:?} / dtype {:?}",
577                mode, dtype
578            ),
579        }
580    }
581
582    fn interpolate_backward(
583        x: FloatTensor<Flex>,
584        grad: FloatTensor<Flex>,
585        output_size: [usize; 2],
586        options: InterpolateOptions,
587    ) -> FloatTensor<Flex> {
588        match (options.mode, x.dtype()) {
589            (InterpolateMode::Nearest, DType::F32) => {
590                interpolate::interpolate_nearest_backward_f32(
591                    x,
592                    grad,
593                    output_size,
594                    options.align_corners,
595                )
596            }
597            (InterpolateMode::Nearest, DType::F64) => {
598                interpolate::interpolate_nearest_backward_f64(
599                    x,
600                    grad,
601                    output_size,
602                    options.align_corners,
603                )
604            }
605            (InterpolateMode::Nearest, DType::F16) => {
606                interpolate::interpolate_nearest_backward_f16(
607                    x,
608                    grad,
609                    output_size,
610                    options.align_corners,
611                )
612            }
613            (InterpolateMode::Nearest, DType::BF16) => {
614                interpolate::interpolate_nearest_backward_bf16(
615                    x,
616                    grad,
617                    output_size,
618                    options.align_corners,
619                )
620            }
621            (InterpolateMode::Bilinear, DType::F32) => {
622                interpolate::interpolate_bilinear_backward_f32(
623                    x,
624                    grad,
625                    output_size,
626                    options.align_corners,
627                )
628            }
629            (InterpolateMode::Bilinear, DType::F64) => {
630                interpolate::interpolate_bilinear_backward_f64(
631                    x,
632                    grad,
633                    output_size,
634                    options.align_corners,
635                )
636            }
637            (InterpolateMode::Bilinear, DType::F16) => {
638                interpolate::interpolate_bilinear_backward_f16(
639                    x,
640                    grad,
641                    output_size,
642                    options.align_corners,
643                )
644            }
645            (InterpolateMode::Bilinear, DType::BF16) => {
646                interpolate::interpolate_bilinear_backward_bf16(
647                    x,
648                    grad,
649                    output_size,
650                    options.align_corners,
651                )
652            }
653            (InterpolateMode::Bicubic, DType::F32) => {
654                interpolate::interpolate_bicubic_backward_f32(
655                    x,
656                    grad,
657                    output_size,
658                    options.align_corners,
659                )
660            }
661            (InterpolateMode::Bicubic, DType::F64) => {
662                interpolate::interpolate_bicubic_backward_f64(
663                    x,
664                    grad,
665                    output_size,
666                    options.align_corners,
667                )
668            }
669            (InterpolateMode::Bicubic, DType::F16) => {
670                interpolate::interpolate_bicubic_backward_f16(
671                    x,
672                    grad,
673                    output_size,
674                    options.align_corners,
675                )
676            }
677            (InterpolateMode::Bicubic, DType::BF16) => {
678                interpolate::interpolate_bicubic_backward_bf16(
679                    x,
680                    grad,
681                    output_size,
682                    options.align_corners,
683                )
684            }
685            (mode, dtype) => {
686                panic!(
687                    "interpolate_backward: unsupported mode {:?} / dtype {:?}",
688                    mode, dtype
689                )
690            }
691        }
692    }
693
694    fn attention(
695        query: FloatTensor<Flex>,
696        key: FloatTensor<Flex>,
697        value: FloatTensor<Flex>,
698        mask: Option<BoolTensor<Flex>>,
699        attn_bias: Option<FloatTensor<Flex>>,
700        options: AttentionModuleOptions,
701    ) -> FloatTensor<Flex> {
702        crate::ops::attention::attention(query, key, value, mask, attn_bias, options)
703    }
704
705    fn rfft(
706        signal: FloatTensor<Flex>,
707        dim: usize,
708        n: Option<usize>,
709    ) -> (FloatTensor<Flex>, FloatTensor<Flex>) {
710        match signal.dtype() {
711            DType::F32 => crate::ops::fft::rfft_f32(signal, dim, n),
712            DType::F64 => crate::ops::fft::rfft_f64(signal, dim, n),
713            DType::F16 => crate::ops::fft::rfft_f16(signal, dim, n),
714            DType::BF16 => crate::ops::fft::rfft_bf16(signal, dim, n),
715            dtype => panic!("rfft: unsupported dtype {:?}", dtype),
716        }
717    }
718
719    fn irfft(
720        spectrum_re: FloatTensor<Flex>,
721        spectrum_im: FloatTensor<Flex>,
722        dim: usize,
723        n: Option<usize>,
724    ) -> FloatTensor<Flex> {
725        match spectrum_re.dtype() {
726            DType::F32 => crate::ops::fft::irfft_f32(spectrum_re, spectrum_im, dim, n),
727            DType::F64 => crate::ops::fft::irfft_f64(spectrum_re, spectrum_im, dim, n),
728            DType::F16 => crate::ops::fft::irfft_f16(spectrum_re, spectrum_im, dim, n),
729            DType::BF16 => crate::ops::fft::irfft_bf16(spectrum_re, spectrum_im, dim, n),
730            dtype => panic!("irfft: unsupported dtype {:?}", dtype),
731        }
732    }
733
734    fn embedding(weights: FloatTensor<Flex>, indices: IntTensor<Flex>) -> FloatTensor<Flex> {
735        let [batch_size, seq_length] = indices.shape().dims();
736        let [_, d_model] = weights.shape().dims();
737
738        let indices = Flex::int_reshape(indices, Shape::from(alloc::vec![batch_size * seq_length]));
739        let output = Flex::float_select(weights, 0, indices);
740        Flex::float_reshape(
741            output,
742            Shape::from(alloc::vec![batch_size, seq_length, d_model]),
743        )
744    }
745
746    fn layer_norm(
747        tensor: FloatTensor<Flex>,
748        gamma: FloatTensor<Flex>,
749        beta: Option<FloatTensor<Flex>>,
750        epsilon: f64,
751    ) -> FloatTensor<Flex> {
752        crate::ops::activation::layer_norm(tensor, gamma, beta, epsilon)
753    }
754
755    fn embedding_backward(
756        weights: FloatTensor<Flex>,
757        output_grad: FloatTensor<Flex>,
758        indices: IntTensor<Flex>,
759    ) -> FloatTensor<Flex> {
760        let [batch_size, seq_length] = indices.shape().dims();
761        let [n_embeddings, d_model] = weights.shape().dims();
762        let dtype = output_grad.dtype();
763
764        let indices = Flex::int_reshape(indices, Shape::from(alloc::vec![batch_size * seq_length]));
765        let output_grad = Flex::float_reshape(
766            output_grad,
767            Shape::from(alloc::vec![batch_size * seq_length, d_model]),
768        );
769        let grad = Flex::float_zeros(
770            Shape::from(alloc::vec![n_embeddings, d_model]),
771            &Default::default(),
772            dtype.into(),
773        );
774        Flex::float_select_add(grad, 0, indices, output_grad)
775    }
776}
burn_flex/ops/module.rs

burn_flex/ops/
module.rs
