Dropout Layer

Problem

Implement a Dropout Layer that randomly zeroes elements of an input tensor during training with probability p, and scales the remaining elements by 1 / (1 - p) (inverted dropout). During evaluation, the input is returned unchanged.

Solution

import numpy as np

def dropout(x: np.ndarray, p: float = 0.5, training: bool = True) -> np.ndarray:
    if not training or p == 0.0:
        return x
    if p == 1.0:
        return np.zeros_like(x)
    mask = (np.random.rand(*x.shape) > p).astype(float)
    return x * mask / (1 - p)

Explanation

1. If not in training mode or p == 0, return the input unchanged.
2. If p == 1, all elements are dropped, so return zeros.
3. Generate a random mask where each element survives with probability 1 - p.
4. Multiply the input by the mask and scale by 1 / (1 - p) so that expected values remain the same at test time (inverted dropout).

Complexity

• Time: O(n) where n is the total number of elements in the tensor
• Space: O(n) for the mask array