pytorch_reader/nested/
adapter.rs

1use super::data::NestedValue;
2
3/// A trait that defines the adapter for a Burn module.
4///
5/// This is used to adapt an incoming module to a Burn module.
6pub trait BurnModuleAdapter: Sized {
7    /// Adapts a module.
8    fn adapt(name: &str, data: NestedValue) -> NestedValue {
9        match name {
10            "BatchNorm" => Self::adapt_batch_norm(data),
11            "Conv1d" => Self::adapt_conv1d(data),
12            "Conv2d" => Self::adapt_conv2d(data),
13            "Conv3d" => Self::adapt_conv3d(data),
14            "ConvTranspose1d" => Self::adapt_conv_transpose_1d(data),
15            "ConvTranspose2d" => Self::adapt_conv_transpose_2d(data),
16            "ConvTranspose3d" => Self::adapt_conv_transpose_3d(data),
17            "Embedding" => Self::adapt_embedding(data),
18            "GroupNorm" => Self::adapt_group_norm(data),
19            "LayerNorm" => Self::adapt_layer_norm(data),
20            "Linear" => Self::adapt_linear(data),
21            _ => data,
22        }
23    }
24
25    /// Adapts a linear module.
26    fn adapt_linear(data: NestedValue) -> NestedValue {
27        data
28    }
29
30    /// Adapts a Convolution 1D module.
31    fn adapt_conv1d(data: NestedValue) -> NestedValue {
32        data
33    }
34
35    /// Adapts a Convolution 2D module.
36    fn adapt_conv2d(data: NestedValue) -> NestedValue {
37        data
38    }
39
40    /// Adapts a Convolution 3D module.
41    fn adapt_conv3d(data: NestedValue) -> NestedValue {
42        data
43    }
44
45    /// Adapts convolution transpose 1D module.
46    fn adapt_conv_transpose_1d(data: NestedValue) -> NestedValue {
47        data
48    }
49
50    /// Adapts convolution transpose 2D module.
51    fn adapt_conv_transpose_2d(data: NestedValue) -> NestedValue {
52        data
53    }
54
55    /// Adapts convolution transpose 2D module.
56    fn adapt_conv_transpose_3d(data: NestedValue) -> NestedValue {
57        data
58    }
59
60    /// Adapts embedding module.
61    fn adapt_embedding(data: NestedValue) -> NestedValue {
62        data
63    }
64
65    /// Adapts group normalization module.
66    fn adapt_group_norm(data: NestedValue) -> NestedValue {
67        data
68    }
69
70    /// Adapts layer normalization module.
71    fn adapt_layer_norm(data: NestedValue) -> NestedValue {
72        data
73    }
74
75    /// Adapts batch normalization module.
76    fn adapt_batch_norm(data: NestedValue) -> NestedValue {
77        data
78    }
79}
80
81/// Default adapter that takes no action.
82pub struct DefaultAdapter;
83impl BurnModuleAdapter for DefaultAdapter {}
pytorch_reader/nested/adapter.rs

pytorch_reader/nested/
adapter.rs
