numcodecs_fourier_network/
modules.rs

1use std::num::NonZeroUsize;
2
3use burn::{
4    config::Config,
5    module::{Module, Param},
6    nn::{BatchNorm, BatchNormConfig, Gelu, Linear, LinearConfig},
7    prelude::Backend,
8    record::{PrecisionSettings, Record},
9    tensor::{Float, Tensor},
10};
11
12#[derive(Debug, Module)]
13pub struct Block<B: Backend> {
14    bn2_1: BatchNorm<B, 0>,
15    gu2_2: Gelu,
16    ln2_3: Linear<B>,
17}
18
19impl<B: Backend> Block<B> {
20    #[expect(clippy::let_and_return)]
21    pub fn forward(&self, x: Tensor<B, 2, Float>) -> Tensor<B, 2, Float> {
22        let x = self.bn2_1.forward(x);
23        let x = self.gu2_2.forward(x);
24        let x = self.ln2_3.forward(x);
25        x
26    }
27}
28
29#[derive(Config, Debug)]
30pub struct BlockConfig {
31    pub fourier_features: NonZeroUsize,
32}
33
34impl BlockConfig {
35    pub fn init<B: Backend>(&self, device: &B::Device) -> Block<B> {
36        Block {
37            bn2_1: BatchNormConfig::new(self.fourier_features.get()).init(device),
38            gu2_2: Gelu,
39            ln2_3: LinearConfig::new(self.fourier_features.get(), self.fourier_features.get())
40                .init(device),
41        }
42    }
43}
44
45#[derive(Debug, Module)]
46pub struct Model<B: Backend> {
47    ln1: Linear<B>,
48    bl2: Vec<Block<B>>,
49    bn3: BatchNorm<B, 0>,
50    gu4: Gelu,
51    ln5: Linear<B>,
52}
53
54impl<B: Backend> Model<B> {
55    #[expect(clippy::let_and_return)]
56    pub fn forward(&self, x: Tensor<B, 2, Float>) -> Tensor<B, 2, Float> {
57        let x = self.ln1.forward(x);
58
59        let mut x = x;
60        for block in &self.bl2 {
61            x = block.forward(x);
62        }
63
64        let x = self.bn3.forward(x);
65        let x = self.gu4.forward(x);
66        let x = self.ln5.forward(x);
67
68        x
69    }
70}
71
72#[derive(Config, Debug)]
73pub struct ModelConfig {
74    pub fourier_features: NonZeroUsize,
75    pub num_blocks: NonZeroUsize,
76}
77
78impl ModelConfig {
79    pub fn init<B: Backend>(&self, device: &B::Device) -> Model<B> {
80        let block = BlockConfig::new(self.fourier_features);
81
82        Model {
83            ln1: LinearConfig::new(self.fourier_features.get() * 2, self.fourier_features.get())
84                .init(device),
85            // (1..num_blocks).into_iter(), fixed in 1.97
86            #[allow(clippy::useless_conversion)]
87            bl2: (1..self.num_blocks.get())
88                .into_iter()
89                .map(|_| block.init(device))
90                .collect(),
91            bn3: BatchNormConfig::new(self.fourier_features.get()).init(device),
92            gu4: Gelu,
93            ln5: LinearConfig::new(self.fourier_features.get(), 1).init(device),
94        }
95    }
96}
97
98pub struct ModelExtra<B: Backend> {
99    pub model: <Model<B> as Module<B>>::Record,
100    pub b_t: Param<Tensor<B, 2, Float>>,
101    pub mean: Param<Tensor<B, 1, Float>>,
102    pub stdv: Param<Tensor<B, 1, Float>>,
103    pub version: crate::FourierNetworkCodecVersion,
104}
105
106impl<B: Backend> Record<B> for ModelExtra<B> {
107    type Item<S: PrecisionSettings> = ModelExtraItem<B, S>;
108
109    fn into_item<S: PrecisionSettings>(self) -> Self::Item<S> {
110        ModelExtraItem {
111            model: self.model.into_item(),
112            b_t: self.b_t.into_item(),
113            mean: self.mean.into_item(),
114            stdv: self.stdv.into_item(),
115            version: self.version,
116        }
117    }
118
119    fn from_item<S: PrecisionSettings>(item: Self::Item<S>, device: &B::Device) -> Self {
120        Self {
121            model: Record::<B>::from_item::<S>(item.model, device),
122            b_t: Record::<B>::from_item::<S>(item.b_t, device),
123            mean: Record::<B>::from_item::<S>(item.mean, device),
124            stdv: Record::<B>::from_item::<S>(item.stdv, device),
125            version: item.version,
126        }
127    }
128}
129
130#[derive(serde::Serialize, serde::Deserialize)]
131#[serde(bound = "")]
132pub struct ModelExtraItem<B: Backend, S: PrecisionSettings> {
133    model: <<Model<B> as Module<B>>::Record as Record<B>>::Item<S>,
134    b_t: <Param<Tensor<B, 2, Float>> as Record<B>>::Item<S>,
135    mean: <Param<Tensor<B, 1, Float>> as Record<B>>::Item<S>,
136    stdv: <Param<Tensor<B, 1, Float>> as Record<B>>::Item<S>,
137    version: crate::FourierNetworkCodecVersion,
138}
numcodecs_fourier_network/modules.rs

numcodecs_fourier_network/
modules.rs
