Tesla ML Coding Interview: 2D Conv Layer Forward in NumPy

Overview

You will implement a 2D convolution forward pass in pure NumPy for channel-first data. The function must accept either a single sample with shape (C_in, H, W) or a batch with shape (N, C_in, H, W), use filters shaped (C_out, C_in, K_h, K_w), optionally add bias (C_out,), and support stride (s_h, s_w) and zero padding (p_h, p_w). The expected output shapes follow the standard conv formula: H_out = floor((H + 2*p_h - K_h) / s_h) + 1 and similarly for W_out.

What you'll be asked (flow)

• Validate inputs (single sample vs batch) and normalize to a batch shape.
• Apply zero padding to spatial axes using NumPy padding or manual slicing.
• Slide each filter across the input with the given stride and compute channel-wise dot products to produce (N, C_out, H_out, W_out).
• Optionally add per-channel bias and return the output in channel-first format.
• Provide unit tests that check small, manually verifiable cases (different strides, paddings, single vs batched, with/without bias).

Skills and signals interviewers look for

You should demonstrate clear mastery of NumPy indexing and slicing, tensor reshaping and broadcasting, shape arithmetic for convolutional output sizes, and writing readable, well-tested code. Bonus signals: a correct vectorized variant (im2col / unfold pattern), micro-benchmarking the optimized vs reference implementations, and a correct parameter-count function using the formula parameters = C_out * C_in * K_h * K_w + (C_out if bias present else 0).

Practical tips

• Normalize inputs early so your core loop assumes a batch dimension.
• Test boundary cases: K > H/W, zero stride mistakes, odd paddings, and non-square kernels.
• When optimizing, focus on removing Python loops (use im2col or NumPy broadcasting) and minimizing temporary allocations.

This prompt is common in ML engineering screens for companies like Tesla and targets practical NumPy competency for convolutional neural network building blocks.