teeny_core/nn/
conv2d.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/// 2-D convolution layer: `output = conv2d(input, weight) [+ b]`
25///
26/// Input shape:  `[N, C_in,  H_in,  W_in ]`
27/// Output shape: `[N, C_out, H_out, W_out]`
28///
29/// where:
30///   `H_out = (H_in + 2*padding_h - kernel_h) / stride_h + 1`
31///   `W_out = (W_in + 2*padding_w - kernel_w) / stride_w + 1`
32///
33/// Type parameters:
34/// - `D`    — element dtype
35/// - `IT`   — input tensor type  (rank 4: `[N, C_in,  H, W]`)
36/// - `OT`   — output tensor type (rank 4: `[N, C_out, H, W]`)
37/// - `RANK` — tensor rank; must be 4 for a valid 2-D convolution
38///
39/// Tensor bounds are on impls, not the struct, so `SymTensor` can have its
40/// own `Layer` impl without a coherence conflict.
41pub struct Conv2d<D: Dtype, IT, OT, const RANK: usize> {
42    /// Number of input channels.
43    pub in_channels: usize,
44    /// Number of output channels.
45    pub out_channels: usize,
46    /// Convolution kernel height.
47    pub kernel_h: usize,
48    /// Convolution kernel width.
49    pub kernel_w: usize,
50    /// Vertical stride.
51    pub stride_h: usize,
52    /// Horizontal stride.
53    pub stride_w: usize,
54    /// Vertical zero-padding.
55    pub padding_h: usize,
56    /// Horizontal zero-padding.
57    pub padding_w: usize,
58    /// Number of blocked/grouped connections (1 = standard convolution).
59    pub groups: usize,
60    /// Whether to add a learned bias.
61    pub has_bias: bool,
62    _pd: PhantomData<(D, IT, OT)>,
63}
64
65impl<D: Dtype, IT, OT, const RANK: usize> Conv2d<D, IT, OT, RANK> {
66    /// Create a new Conv2d layer.
67    ///
68    /// - `kernel_size` — `(height, width)` of the convolution kernel
69    /// - `stride`      — `(height, width)` step between kernel applications
70    /// - `padding`     — `(height, width)` zero-padding added to each spatial side
71    pub fn new(
72        in_channels: usize,
73        out_channels: usize,
74        kernel_size: (usize, usize),
75        stride: (usize, usize),
76        padding: (usize, usize),
77        has_bias: bool,
78    ) -> Self {
79        Self {
80            in_channels,
81            out_channels,
82            kernel_h: kernel_size.0,
83            kernel_w: kernel_size.1,
84            stride_h: stride.0,
85            stride_w: stride.1,
86            padding_h: padding.0,
87            padding_w: padding.1,
88            groups: 1,
89            has_bias,
90            _pd: PhantomData,
91        }
92    }
93
94    /// Creates a new grouped `Conv2d` layer (see [`Conv2d::groups`]).
95    pub fn new_grouped(
96        in_channels: usize,
97        out_channels: usize,
98        kernel_size: (usize, usize),
99        stride: (usize, usize),
100        padding: (usize, usize),
101        has_bias: bool,
102        groups: usize,
103    ) -> Self {
104        Self {
105            in_channels,
106            out_channels,
107            kernel_h: kernel_size.0,
108            kernel_w: kernel_size.1,
109            stride_h: stride.0,
110            stride_w: stride.1,
111            padding_h: padding.0,
112            padding_w: padding.1,
113            groups,
114            has_bias,
115            _pd: PhantomData,
116        }
117    }
118}
119
120impl<D: Dtype, IT: Tensor<D, RANK> + EagerTensor, OT: Tensor<D, RANK>, const RANK: usize> Layer<IT>
121    for Conv2d<D, IT, OT, RANK>
122{
123    type Output = OT;
124
125    fn call(&self, _input: IT) -> Self::Output {
126        todo!()
127    }
128}
teeny_core/nn/conv2d.rs

teeny_core/nn/
conv2d.rs
