zyx-nn 0.15.3

Zyx nn modules
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66

// Copyright (C) 2025 zk4x
// SPDX-License-Identifier: LGPL-3.0-only

use zyx::{DType, IntoShape, Tensor, ZyxError};
use zyx_derive::Module;

/// Applies a 2D convolution over an input signal composed of several input planes.
///
/// See: https://pytorch.org/docs/stable/generated/torch.nn.Conv2d
#[derive(Debug, Module)]
#[cfg_attr(feature = "py", pyo3::pyclass)]
pub struct Conv2d {
    stride: Vec<u64>,
    dilation: Vec<u64>,
    groups: u64,
    padding: Vec<u64>,
    /// weight
    pub weight: Tensor,
    /// bias
    pub bias: Option<Tensor>,
}

impl Conv2d {
    /// Initialize Conv2d
    pub fn new(
        in_channels: u64,
        out_channels: u64,
        kernel_size: impl IntoShape,
        stride: impl IntoShape,
        padding: impl IntoShape,
        dilation: impl IntoShape,
        groups: u64,
        bias: bool,
        dtype: DType,
    ) -> Result<Self, ZyxError> {
        let mut kernel_size: Vec<u64> = kernel_size.into_shape().collect();
        kernel_size.push(2);
        let scale = 1f32 / ((in_channels * kernel_size.iter().product::<u64>()) as f32).sqrt();
        let mut weight_shape = vec![out_channels, in_channels / groups];
        weight_shape.extend(kernel_size);
        Ok(Conv2d {
            stride: stride.into_shape().collect(),
            dilation: dilation.into_shape().collect(),
            groups,
            padding: padding.into_shape().collect(),
            weight: Tensor::uniform(weight_shape, -scale..scale)?.cast(dtype),
            bias: if bias {
                Some(Tensor::uniform(out_channels, -scale..scale)?.cast(dtype))
            } else {
                None
            },
        })
    }

    /// Forward conv2d layer
    pub fn forward(&self, x: impl Into<Tensor>) -> Result<Tensor, ZyxError> {
        x.into().conv(
            &self.weight,
            self.bias.as_ref(),
            self.groups,
            &self.stride,
            &self.dilation,
            &self.padding,
        )
    }
}