teeny_core/nn/
mod.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
17/// Activation layers (ReLU, GELU, sigmoid, etc).
18pub mod activation;
19/// Batch normalization.
20pub mod batchnorm;
21/// 1-D convolution.
22pub mod conv1d;
23/// 2-D convolution.
24pub mod conv2d;
25/// 3-D convolution.
26pub mod conv3d;
27/// Flattening a tensor to fewer dimensions.
28pub mod flatten;
29/// Group normalization.
30pub mod groupnorm;
31/// Instance normalization.
32pub mod instancenorm;
33/// Layer normalization.
34pub mod layernorm;
35/// Fully-connected (dense) layers.
36pub mod linear;
37/// Padding layers.
38pub mod pad;
39/// Pooling layers.
40pub mod pool;
41/// RMS normalization.
42pub mod rmsnorm;
43
44/// A neural network layer: something callable on an input `I`, producing `Output`.
45pub trait Layer<I> {
46    /// This layer's output type.
47    type Output;
48
49    /// Applies this layer to `input`.
50    fn call(&self, input: I) -> Self::Output;
51}
52
53impl<F, I, O> Layer<I> for F
54where
55    F: Fn(I) -> O,
56{
57    type Output = O;
58
59    fn call(&self, input: I) -> O {
60        self(input)
61    }
62}
63
64#[cfg(test)]
65mod tests {
66    use crate::{
67        dtype::{Dtype, EagerTensor, RankedTensor, Tensor},
68        nn::{
69            Layer,
70            activation::{relu::Relu, softmax::Softmax},
71            linear::Linear,
72        },
73        sequential,
74    };
75
76    /// Dummy tensor: satisfies `Tensor<D, RANK>` for any dtype and rank.
77    /// Shape is all-zeros (dynamic/unknown). Used only to verify type composition.
78    #[derive(Copy, Clone)]
79    struct DummyTensor;
80
81    impl<D: Dtype, const RANK: usize> RankedTensor<D, RANK> for DummyTensor {
82        const SHAPE: [usize; RANK] = [0; RANK];
83    }
84
85    impl<D: Dtype, const RANK: usize> Tensor<D, RANK> for DummyTensor {}
86    impl EagerTensor for DummyTensor {}
87
88    /// A 2-layer MLP: Linear(784→128) → ReLU → Linear(128→10) → Softmax
89    ///
90    /// RANK=2 throughout so a batch dimension `[batch, features]` flows
91    /// through every layer without changing rank.
92    #[test]
93    fn test_mlp_composition() {
94        let _model = sequential![
95            Linear::<f32, DummyTensor, DummyTensor, 2>::new(784, 128, true),
96            Relu::<f32, DummyTensor, 2>::new(),
97            Linear::<f32, DummyTensor, DummyTensor, 2>::new(128, 10, true),
98            Softmax::<f32, DummyTensor, 2>::new(1)
99        ];
100    }
101}
teeny_core/nn/mod.rs

teeny_core/nn/
mod.rs
