teeny_core/nn/
instancenorm.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// ─── InstanceNorm1d ──────────────────────────────────────────────────────────
25
26/// Instance normalisation over a 3-D `[N, C, L]` input.
27///
28/// Normalises each (N, C) pair independently over the L spatial dimension.
29/// Output shape equals input shape.
30///
31/// Parameters:
32/// - `num_features`        — number of channels C
33/// - `eps`                 — numerical stability constant (default 1e-5)
34/// - `momentum`            — EMA weight for running stats (default 0.1)
35/// - `affine`              — if true, learns per-channel γ and β (default false)
36/// - `track_running_stats` — maintain running mean/var (default false)
37pub struct InstanceNorm1d<D: Dtype, IT, OT, const RANK: usize> {
38    /// Number of channels/features.
39    pub num_features: usize,
40    /// Numerical stability constant added to the variance.
41    pub eps: f64,
42    /// Weight of the current batch in the running-stats exponential moving average.
43    pub momentum: f64,
44    /// Whether to learn per-channel scale (gamma) and shift (beta) parameters.
45    pub affine: bool,
46    /// Whether to maintain running mean/variance across batches.
47    pub track_running_stats: bool,
48    _pd: PhantomData<(D, IT, OT)>,
49}
50
51impl<D: Dtype, IT, OT, const RANK: usize> InstanceNorm1d<D, IT, OT, RANK> {
52    /// Creates a new layer with default `eps`/`momentum`, affine/running-stats off.
53    pub fn new(num_features: usize) -> Self {
54        Self {
55            num_features,
56            eps: 1e-5,
57            momentum: 0.1,
58            affine: false,
59            track_running_stats: false,
60            _pd: PhantomData,
61        }
62    }
63
64    /// Sets the numerical stability constant.
65    pub fn with_eps(mut self, eps: f64) -> Self {
66        self.eps = eps;
67        self
68    }
69    /// Sets the running-stats EMA momentum.
70    pub fn with_momentum(mut self, m: f64) -> Self {
71        self.momentum = m;
72        self
73    }
74    /// Sets whether to learn per-channel scale/shift parameters.
75    pub fn with_affine(mut self, a: bool) -> Self {
76        self.affine = a;
77        self
78    }
79    /// Sets whether to maintain running mean/variance across batches.
80    pub fn with_track_running_stats(mut self, t: bool) -> Self {
81        self.track_running_stats = t;
82        self
83    }
84}
85
86impl<D: Dtype, IT: Tensor<D, RANK> + EagerTensor, OT: Tensor<D, RANK>, const RANK: usize> Layer<IT>
87    for InstanceNorm1d<D, IT, OT, RANK>
88{
89    type Output = OT;
90    fn call(&self, _input: IT) -> Self::Output {
91        todo!()
92    }
93}
94
95// ─── InstanceNorm2d ──────────────────────────────────────────────────────────
96
97/// Instance normalisation over a 4-D `[N, C, H, W]` input.
98pub struct InstanceNorm2d<D: Dtype, IT, OT, const RANK: usize> {
99    /// Number of channels/features.
100    pub num_features: usize,
101    /// Numerical stability constant added to the variance.
102    pub eps: f64,
103    /// Weight of the current batch in the running-stats exponential moving average.
104    pub momentum: f64,
105    /// Whether to learn per-channel scale (gamma) and shift (beta) parameters.
106    pub affine: bool,
107    /// Whether to maintain running mean/variance across batches.
108    pub track_running_stats: bool,
109    _pd: PhantomData<(D, IT, OT)>,
110}
111
112impl<D: Dtype, IT, OT, const RANK: usize> InstanceNorm2d<D, IT, OT, RANK> {
113    /// Creates a new layer with default `eps`/`momentum`, affine/running-stats off.
114    pub fn new(num_features: usize) -> Self {
115        Self {
116            num_features,
117            eps: 1e-5,
118            momentum: 0.1,
119            affine: false,
120            track_running_stats: false,
121            _pd: PhantomData,
122        }
123    }
124
125    /// Sets the numerical stability constant.
126    pub fn with_eps(mut self, eps: f64) -> Self {
127        self.eps = eps;
128        self
129    }
130    /// Sets the running-stats EMA momentum.
131    pub fn with_momentum(mut self, m: f64) -> Self {
132        self.momentum = m;
133        self
134    }
135    /// Sets whether to learn per-channel scale/shift parameters.
136    pub fn with_affine(mut self, a: bool) -> Self {
137        self.affine = a;
138        self
139    }
140    /// Sets whether to maintain running mean/variance across batches.
141    pub fn with_track_running_stats(mut self, t: bool) -> Self {
142        self.track_running_stats = t;
143        self
144    }
145}
146
147impl<D: Dtype, IT: Tensor<D, RANK> + EagerTensor, OT: Tensor<D, RANK>, const RANK: usize> Layer<IT>
148    for InstanceNorm2d<D, IT, OT, RANK>
149{
150    type Output = OT;
151    fn call(&self, _input: IT) -> Self::Output {
152        todo!()
153    }
154}
155
156// ─── InstanceNorm3d ──────────────────────────────────────────────────────────
157
158/// Instance normalisation over a 5-D `[N, C, D, H, W]` input.
159pub struct InstanceNorm3d<D: Dtype, IT, OT, const RANK: usize> {
160    /// Number of channels/features.
161    pub num_features: usize,
162    /// Numerical stability constant added to the variance.
163    pub eps: f64,
164    /// Weight of the current batch in the running-stats exponential moving average.
165    pub momentum: f64,
166    /// Whether to learn per-channel scale (gamma) and shift (beta) parameters.
167    pub affine: bool,
168    /// Whether to maintain running mean/variance across batches.
169    pub track_running_stats: bool,
170    _pd: PhantomData<(D, IT, OT)>,
171}
172
173impl<D: Dtype, IT, OT, const RANK: usize> InstanceNorm3d<D, IT, OT, RANK> {
174    /// Creates a new layer with default `eps`/`momentum`, affine/running-stats off.
175    pub fn new(num_features: usize) -> Self {
176        Self {
177            num_features,
178            eps: 1e-5,
179            momentum: 0.1,
180            affine: false,
181            track_running_stats: false,
182            _pd: PhantomData,
183        }
184    }
185
186    /// Sets the numerical stability constant.
187    pub fn with_eps(mut self, eps: f64) -> Self {
188        self.eps = eps;
189        self
190    }
191    /// Sets the running-stats EMA momentum.
192    pub fn with_momentum(mut self, m: f64) -> Self {
193        self.momentum = m;
194        self
195    }
196    /// Sets whether to learn per-channel scale/shift parameters.
197    pub fn with_affine(mut self, a: bool) -> Self {
198        self.affine = a;
199        self
200    }
201    /// Sets whether to maintain running mean/variance across batches.
202    pub fn with_track_running_stats(mut self, t: bool) -> Self {
203        self.track_running_stats = t;
204        self
205    }
206}
207
208impl<D: Dtype, IT: Tensor<D, RANK> + EagerTensor, OT: Tensor<D, RANK>, const RANK: usize> Layer<IT>
209    for InstanceNorm3d<D, IT, OT, RANK>
210{
211    type Output = OT;
212    fn call(&self, _input: IT) -> Self::Output {
213        todo!()
214    }
215}
teeny_core/nn/instancenorm.rs

teeny_core/nn/
instancenorm.rs
