Katkat

Implement 2D convolution forward pass

TeslaMachine Learning Engineer · Technical Screen

Last updated

This question evaluates understanding of 2D convolution mechanics and practical low-level tensor manipulation, including kernel shape handling, padding and stride effects, multi-channel input/output interactions, and optional bias integration.

Problem

Implement the forward pass of a 2D convolution (conv2d) from scratch (no deep learning libraries).

You are given:

  • Input tensor x with shape (N, C_in, H, W) (NCHW layout)
  • Filter weights w with shape (C_out, C_in, K_h, K_w)
  • Optional bias b with shape (C_out,) (may be None)
  • Integer stride s_h, s_w
  • Integer padding p_h, p_w (zero-padding applied to height/width)

Compute the output tensor y with shape (N, C_out, H_out, W_out) where:

[ H_{out} = \left\lfloor \frac{H + 2p_h - K_h}{s_h} \right\rfloor + 1, \quad W_{out} = \left\lfloor \frac{W + 2p_w - K_w}{s_w} \right\rfloor + 1 ]

For each output element:

[ y[n, c_{out}, i, j] = \sum_{c_{in}=0}^{C_{in}-1} \sum_{u=0}^{K_h-1} \sum_{v=0}^{K_w-1} ; x_{pad}[n, c_{in}, i\cdot s_h + u, j\cdot s_w + v] \cdot w[c_{out}, c_{in}, u, v] + b[c_{out}] ]

where x_pad is x padded with zeros by (p_h, p_w).

Requirements

  • Return y as a dense numeric array/tensor.
  • Do not use existing convolution operators.
  • Handle edge cases such as b is None, non-square kernels, and different strides/padding.

Constraints (typical for an interview unit test)

  • Sizes are small enough that an (O(N\cdot C_{out}\cdot C_{in}\cdot H_{out}\cdot W_{out}\cdot K_h\cdot K_w)) implementation passes.
  • Inputs are floating point numbers.

Example (shape check)

If x is (1, 3, 32, 32), w is (8, 3, 3, 3), stride (1,1), padding (1,1), then output shape is (1, 8, 32, 32).