Katkat
Back

Implement common neural network layers in PyTorch

Startups.ComMachine Learning Engineer · Onsite
MediumMachine Learning EngineerOnsiteCoding & Algorithms

Last updated

This question evaluates proficiency in implementing core neural-network primitives (ReLU, numerically-stable Softmax, LayerNorm, BatchNorm with running statistics, and RMSNorm), numerical stability and normalization concepts, and understanding of autograd-compatible tensor operations.

Implement the following neural-network building blocks from scratch in PyTorch (using only basic tensor ops such as matmul, sum, mean, var, exp, max, reshape, transpose, broadcasting). Do not call torch.nn.* or torch.nn.functional.* implementations of the same ops.

  1. ReLU
  • Input: tensor x of any shape.
  • Output: y = max(x, 0).
  1. Softmax (numerically stable)
  • Input: tensor x and an integer dim.
  • Output: softmax along dim.
  • Must be stable for large/small values.
  1. LayerNorm
  • Input: x with shape [*, normalized_dim] (or more generally normalize over the last k dims), learnable gamma, beta, and eps.
  • Output: normalized tensor.
  1. BatchNorm (training and inference)
  • Input: x (e.g., [N, C, H, W]), learnable gamma, beta, eps, momentum, plus running_mean and running_var.
  • Implement:
  • Training forward: compute batch stats and update running stats.
  • Inference forward: use running stats.
  1. RMSNorm
  • Input: x, learnable gamma, eps.
  • Normalize by RMS (no mean subtraction).

For each component, describe how you would verify correctness (e.g., compare to a reference implementation, run gradient checks / finite-difference checks, test edge cases such as fp16, tiny variance, and different tensor shapes).