Skip to main content

teeny_core/nn/
pad.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
24macro_rules! pad_layer {
25    (
26        #[doc = $doc:expr]
27        $name:ident,
28        $( $field:ident : $fty:ty ),+
29        $(,)?
30    ) => {
31        #[doc = $doc]
32        pub struct $name<D: Dtype, IT, OT, const RANK: usize> {
33            $( #[doc = "Padding/config parameter; see the struct docs."] pub $field: $fty, )+
34            _pd: PhantomData<(D, IT, OT)>,
35        }
36
37        impl<D: Dtype, IT, OT, const RANK: usize> $name<D, IT, OT, RANK> {
38            #[doc = "Creates a new layer with the given padding/config parameters."]
39            pub fn new( $( $field: $fty ),+ ) -> Self {
40                Self { $( $field, )+ _pd: PhantomData }
41            }
42        }
43
44        impl<D: Dtype, IT: Tensor<D, RANK> + EagerTensor, OT: Tensor<D, RANK>, const RANK: usize>
45            Layer<IT> for $name<D, IT, OT, RANK>
46        {
47            type Output = OT;
48            fn call(&self, _input: IT) -> Self::Output {
49                todo!()
50            }
51        }
52    };
53}
54
55// ConstantPad — adds a constant fill value on each side
56pad_layer!(#[doc = "1-D constant padding: pads with a fixed value on each side."] ConstantPad1d, pad_left: usize, pad_right: usize, value: f64);
57pad_layer!(#[doc = "2-D constant padding: pads with a fixed value on each side."] ConstantPad2d, pad_l: usize, pad_r: usize, pad_t: usize, pad_b: usize, value: f64);
58pad_layer!(#[doc = "3-D constant padding: pads with a fixed value on each side."] ConstantPad3d, pad_d1: usize, pad_d2: usize, pad_h1: usize, pad_h2: usize, pad_w1: usize, pad_w2: usize, value: f64);
59
60// ReflectionPad — pads by reflecting the input at the boundary
61pad_layer!(#[doc = "1-D reflection padding: pads by reflecting the input at the boundary."] ReflectionPad1d, pad_left: usize, pad_right: usize);
62pad_layer!(#[doc = "2-D reflection padding: pads by reflecting the input at the boundary."] ReflectionPad2d, pad_l: usize, pad_r: usize, pad_t: usize, pad_b: usize);
63pad_layer!(#[doc = "3-D reflection padding: pads by reflecting the input at the boundary."] ReflectionPad3d, pad_d1: usize, pad_d2: usize, pad_h1: usize, pad_h2: usize, pad_w1: usize, pad_w2: usize);
64
65// ReplicationPad — pads by replicating the edge values
66pad_layer!(#[doc = "1-D replication padding: pads by replicating edge values."] ReplicationPad1d, pad_left: usize, pad_right: usize);
67pad_layer!(#[doc = "2-D replication padding: pads by replicating edge values."] ReplicationPad2d, pad_l: usize, pad_r: usize, pad_t: usize, pad_b: usize);
68pad_layer!(#[doc = "3-D replication padding: pads by replicating edge values."] ReplicationPad3d, pad_d1: usize, pad_d2: usize, pad_h1: usize, pad_h2: usize, pad_w1: usize, pad_w2: usize);
69
70// CircularPad — pads by wrapping around (circular/periodic boundary)
71pad_layer!(#[doc = "1-D circular padding: pads by wrapping around (periodic boundary)."] CircularPad1d, pad_left: usize, pad_right: usize);
72pad_layer!(#[doc = "2-D circular padding: pads by wrapping around (periodic boundary)."] CircularPad2d, pad_l: usize, pad_r: usize, pad_t: usize, pad_b: usize);
73pad_layer!(#[doc = "3-D circular padding: pads by wrapping around (periodic boundary)."] CircularPad3d, pad_d1: usize, pad_d2: usize, pad_h1: usize, pad_h2: usize, pad_w1: usize, pad_w2: usize);