ReLU Forward and Backward

Implement the ReLU (Rectified Linear Unit) activation, the most common nonlinearity in deep networks, in both directions. The forward pass clamps negatives to zero; the backward pass routes the upstream gradient through only the units that were active. Interviewers use this to check that you can turn a piecewise function into vectorized NumPy and that you gate the gradient on the layer’s input rather than its output.

Definition

ReLU is applied elementwise:

relu(x) = max(0, x)

Its derivative is the indicator that the input was positive:

d relu / d x = 1 if x > 0 else 0     (undefined at x == 0; use 0 by convention)

By the chain rule, the gradient flowing back to the input is the upstream gradient multiplied by that local derivative:

dX = dOut * (x > 0)

Task

Complete two functions, fully vectorized with NO Python loops and NO deep-learning framework (NumPy only):

• relu(x) returns max(0, x) applied elementwise, same shape as x.
• relu_backward(dOut, x) returns the gradient with respect to x: it passes dOut through where x > 0 and returns 0 where x <= 0. Gate on x (the forward input), not on the forward output.

Example

x = np.array([[-1.0, 0.0, 2.0],
              [ 3.0, -4.0, 5.0]])
relu(x)
# -> array([[0., 0., 2.],
#           [3., 0., 5.]])

dOut = np.array([[0.1, 0.2, 0.3],
                 [0.4, 0.5, 0.6]])
relu_backward(dOut, x)
# -> array([[0. , 0. , 0.3],
#           [0.4, 0. , 0.6]])

The gradient at position (0, 2) survives as 0.3 because x = 2.0 > 0; the gradient at (0, 0) is zeroed because x = -1.0, and at (0, 1) it is zeroed because x = 0.0 is not strictly positive.

Constraints

• x and dOut are float arrays of identical shape (any dimensionality); use vectorized NumPy, no Python loop over elements.
• At x == 0 the derivative is taken to be 0 (gate on x > 0, strictly greater).
• Do not modify dOut or x in place; return a new array.

Approach

The forward pass is a single elementwise clamp: np.maximum(0, x) compares each element against zero and keeps the larger, which is exactly max(0, x). The backward pass is a gated pass-through: ReLU’s local derivative is 1 for active units (x > 0) and 0 otherwise, so by the chain rule the input gradient is dOut masked by that boolean, i.e. dOut * (x > 0). Both are pure NumPy vector ops with no loops, and both allocate new arrays so the inputs are left untouched.

Solution

import numpy as np

def relu(x: np.ndarray) -> np.ndarray:
    # elementwise max(0, x); np.maximum broadcasts the scalar 0 across x
    return np.maximum(0.0, x)

def relu_backward(dOut: np.ndarray, x: np.ndarray) -> np.ndarray:
    # local derivative is 1 where x > 0, else 0; gate the upstream gradient with it
    mask = (x > 0)                      # boolean array, same shape as x
    return dOut * mask                  # False -> 0.0, True -> dOut, new array

Walkthrough

On the example x = [[-1, 0, 2], [3, -4, 5]]:

• relu(x): np.maximum(0, x) compares each element to 0 and keeps the larger, giving [[0, 0, 2], [3, 0, 5]]. The negatives -1 and -4 become 0, the 0 stays 0, and the positives pass through.
• relu_backward(dOut, x) with dOut = [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]:
  • mask = (x > 0) = [[False, False, True], [True, False, True]] — note 0 is False because the test is strictly greater.
  • dOut * mask multiplies each gradient by 1.0 or 0.0, yielding [[0, 0, 0.3], [0.4, 0, 0.6]].

The 0.3 and 0.4/0.6 gradients survive exactly at the active units; every other position is zeroed.

Complexity & notes

• Time O(N) and space O(N) for both functions, where N is the number of elements: one pass to compute the elementwise result and one output array. No matmul, no reduction.
• Gate on the forward INPUT x, not the output. Gating on the output relu(x) > 0 gives the same mask here, but keeping the original input is the standard cache and generalizes to activations whose output does not reveal the input sign.
• The strict x > 0 encodes the x == 0 convention: the derivative is undefined at the kink, and 0 is the near-universal choice (PyTorch and TensorFlow do the same). Using >= would leak gradient through dead-zero units.
• dOut * mask relies on NumPy casting bool to 0.0/1.0 during multiplication, so the result stays float and no branch is needed. It returns a fresh array, so dOut and x are not mutated.
• In a full layer you would cache x (or the mask) during relu and reuse it in relu_backward; here x is passed back in explicitly to keep the two functions independent and hand-checkable.