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

burn_flex/ops/
module.rs
