teeny_core/nn/batchnorm.rs
1/*
2 * Copyright (c) 2026 Teenygrad.
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 * http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17use core::marker::PhantomData;
18
19use crate::{
20 dtype::{Dtype, EagerTensor, Tensor},
21 nn::Layer,
22};
23
24// ─── BatchNorm1d ─────────────────────────────────────────────────────────────
25
26/// Batch normalisation over a 2-D `[N, C]` or 3-D `[N, C, L]` input.
27///
28/// During training the per-batch mean and variance are used; the running
29/// statistics are updated with exponential moving average. During inference
30/// the frozen running statistics are used instead.
31///
32/// Parameters:
33/// - `num_features` — number of channels C
34/// - `eps` — numerical stability constant added to the variance
35/// - `momentum` — weight of the current batch in the running-stats EMA
36/// - `affine` — if true, learns per-channel γ (weight) and β (bias)
37/// - `track_running_stats` — if true, maintains running_mean / running_var
38pub struct BatchNorm1d<D: Dtype, IT, OT, const RANK: usize> {
39 /// Number of channels/features.
40 pub num_features: usize,
41 /// Numerical stability constant added to the variance.
42 pub eps: f64,
43 /// Weight of the current batch in the running-stats exponential moving average.
44 pub momentum: f64,
45 /// Whether to learn per-channel scale (gamma) and shift (beta) parameters.
46 pub affine: bool,
47 /// Whether to maintain running mean/variance across batches.
48 pub track_running_stats: bool,
49 _pd: PhantomData<(D, IT, OT)>,
50}
51
52impl<D: Dtype, IT, OT, const RANK: usize> BatchNorm1d<D, IT, OT, RANK> {
53 /// Creates a new layer with default `eps`/`momentum`/`affine`/`track_running_stats`.
54 pub fn new(num_features: usize) -> Self {
55 Self {
56 num_features,
57 eps: 1e-5,
58 momentum: 0.1,
59 affine: true,
60 track_running_stats: true,
61 _pd: PhantomData,
62 }
63 }
64
65 /// Sets the numerical stability constant.
66 pub fn with_eps(mut self, eps: f64) -> Self {
67 self.eps = eps;
68 self
69 }
70
71 /// Sets the running-stats EMA momentum.
72 pub fn with_momentum(mut self, momentum: f64) -> Self {
73 self.momentum = momentum;
74 self
75 }
76
77 /// Sets whether to learn per-channel scale/shift parameters.
78 pub fn with_affine(mut self, affine: bool) -> Self {
79 self.affine = affine;
80 self
81 }
82
83 /// Sets whether to maintain running mean/variance across batches.
84 pub fn with_track_running_stats(mut self, track: bool) -> Self {
85 self.track_running_stats = track;
86 self
87 }
88}
89
90impl<D: Dtype, IT: Tensor<D, RANK> + EagerTensor, OT: Tensor<D, RANK>, const RANK: usize> Layer<IT>
91 for BatchNorm1d<D, IT, OT, RANK>
92{
93 type Output = OT;
94
95 fn call(&self, _input: IT) -> Self::Output {
96 todo!()
97 }
98}
99
100// ─── BatchNorm2d ─────────────────────────────────────────────────────────────
101
102/// Batch normalisation over a 4-D `[N, C, H, W]` input.
103///
104/// Applies the same statistics as `BatchNorm1d` but treats each (C) channel
105/// independently across the (N, H, W) axes.
106pub struct BatchNorm2d<D: Dtype, IT, OT, const RANK: usize> {
107 /// Number of channels/features.
108 pub num_features: usize,
109 /// Numerical stability constant added to the variance.
110 pub eps: f64,
111 /// Weight of the current batch in the running-stats exponential moving average.
112 pub momentum: f64,
113 /// Whether to learn per-channel scale (gamma) and shift (beta) parameters.
114 pub affine: bool,
115 /// Whether to maintain running mean/variance across batches.
116 pub track_running_stats: bool,
117 _pd: PhantomData<(D, IT, OT)>,
118}
119
120impl<D: Dtype, IT, OT, const RANK: usize> BatchNorm2d<D, IT, OT, RANK> {
121 /// Creates a new layer with default `eps`/`momentum`/`affine`/`track_running_stats`.
122 pub fn new(num_features: usize) -> Self {
123 Self {
124 num_features,
125 eps: 1e-5,
126 momentum: 0.1,
127 affine: true,
128 track_running_stats: true,
129 _pd: PhantomData,
130 }
131 }
132
133 /// Sets the numerical stability constant.
134 pub fn with_eps(mut self, eps: f64) -> Self {
135 self.eps = eps;
136 self
137 }
138
139 /// Sets the running-stats EMA momentum.
140 pub fn with_momentum(mut self, momentum: f64) -> Self {
141 self.momentum = momentum;
142 self
143 }
144
145 /// Sets whether to learn per-channel scale/shift parameters.
146 pub fn with_affine(mut self, affine: bool) -> Self {
147 self.affine = affine;
148 self
149 }
150
151 /// Sets whether to maintain running mean/variance across batches.
152 pub fn with_track_running_stats(mut self, track: bool) -> Self {
153 self.track_running_stats = track;
154 self
155 }
156}
157
158impl<D: Dtype, IT: Tensor<D, RANK> + EagerTensor, OT: Tensor<D, RANK>, const RANK: usize> Layer<IT>
159 for BatchNorm2d<D, IT, OT, RANK>
160{
161 type Output = OT;
162
163 fn call(&self, _input: IT) -> Self::Output {
164 todo!()
165 }
166}
167
168// ─── BatchNorm3d ─────────────────────────────────────────────────────────────
169
170/// Batch normalisation over a 5-D `[N, C, D, H, W]` input.
171///
172/// Applies the same statistics as `BatchNorm1d` but treats each (C) channel
173/// independently across the (N, D, H, W) axes.
174pub struct BatchNorm3d<D: Dtype, IT, OT, const RANK: usize> {
175 /// Number of channels/features.
176 pub num_features: usize,
177 /// Numerical stability constant added to the variance.
178 pub eps: f64,
179 /// Weight of the current batch in the running-stats exponential moving average.
180 pub momentum: f64,
181 /// Whether to learn per-channel scale (gamma) and shift (beta) parameters.
182 pub affine: bool,
183 /// Whether to maintain running mean/variance across batches.
184 pub track_running_stats: bool,
185 _pd: PhantomData<(D, IT, OT)>,
186}
187
188impl<D: Dtype, IT, OT, const RANK: usize> BatchNorm3d<D, IT, OT, RANK> {
189 /// Creates a new layer with default `eps`/`momentum`/`affine`/`track_running_stats`.
190 pub fn new(num_features: usize) -> Self {
191 Self {
192 num_features,
193 eps: 1e-5,
194 momentum: 0.1,
195 affine: true,
196 track_running_stats: true,
197 _pd: PhantomData,
198 }
199 }
200
201 /// Sets the numerical stability constant.
202 pub fn with_eps(mut self, eps: f64) -> Self {
203 self.eps = eps;
204 self
205 }
206
207 /// Sets the running-stats EMA momentum.
208 pub fn with_momentum(mut self, momentum: f64) -> Self {
209 self.momentum = momentum;
210 self
211 }
212
213 /// Sets whether to learn per-channel scale/shift parameters.
214 pub fn with_affine(mut self, affine: bool) -> Self {
215 self.affine = affine;
216 self
217 }
218
219 /// Sets whether to maintain running mean/variance across batches.
220 pub fn with_track_running_stats(mut self, track: bool) -> Self {
221 self.track_running_stats = track;
222 self
223 }
224}
225
226impl<D: Dtype, IT: Tensor<D, RANK> + EagerTensor, OT: Tensor<D, RANK>, const RANK: usize> Layer<IT>
227 for BatchNorm3d<D, IT, OT, RANK>
228{
229 type Output = OT;
230
231 fn call(&self, _input: IT) -> Self::Output {
232 todo!()
233 }
234}