Struct Conv2d 

pub struct Conv2d {
    pub weight: Parameter,
    pub bias: Option<Parameter>,
    pub in_channels: usize,
    pub out_channels: usize,
    pub kernel_size: usize,
    pub stride: usize,
    pub padding: usize,
}

2D convolutional layer.

Uses the im2col algorithm: reshape input patches into columns, then compute the convolution as a single matmul per batch element.

Weight shape: [out_channels, in_channels * kernel_size * kernel_size]. Bias shape: [out_channels] (optional).

Forward: for each batch element:

1. x_col = input[b].im2col(k, stride, padding) → [col_height, num_patches]
2. out_b = weight @ x_col → [out_channels, num_patches]
3. Stack all out_b → [batch, out_channels, num_patches]
4. Add channel bias if present.
5. Reshape to [batch, out_channels, out_h, out_w].

All ops are individually tape-recorded — the autograd engine handles the entire backward pass via Kahn’s traversal.

Fields

weight: Parameter

bias: Option<Parameter>

in_channels: usize

out_channels: usize

kernel_size: usize

stride: usize

padding: usize

Implementations

impl Conv2d

pub fn new( in_channels: usize, out_channels: usize, kernel_size: usize, stride: usize, padding: usize, with_bias: bool, ) -> Self

Create a new Conv2d layer with Kaiming Uniform initialization.

pub fn forward(&self, input: &Tensor) -> Tensor

Forward pass.

input shape: [batch, in_channels, height, width]. Output shape: [batch, out_channels, out_h, out_w].

Each step is a tracked op so the autograd tape records the full graph: slice_batch → im2col → matmul → add_channel_bias → stack.

Trait Implementations

impl Module for Conv2d

fn parameters(&self) -> Vec<Parameter>

Return clones of all learnable parameters.

fn state_dict(&self, prefix: &str) -> HashMap<String, Tensor>

Serialize all parameters into a flat name → Tensor map.

fn load_state_dict( &mut self, dict: &HashMap<String, Tensor>, prefix: &str, ) -> Result<(), AutogradError>

Load parameters from a flat name → Tensor map.

fn train(&mut self)

Switch to training mode.

fn eval(&mut self)

Switch to evaluation mode.