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

burn_flex/ops/
module.rs
