Medium Primitives

Implement Dropout

Dropout is a training-time regularization technique that randomly drops some activations. In inverted dropout, the values that remain are scaled up by the reciprocal of their keep probability.

This exercise supplies the mask directly. Your function applies that mask and scaling; it must not sample random values or create a new mask.

Your task

Implement dropout(x, mask, p=0.5) using basic elementwise operations. Do not call a built-in dropout layer or function.

x is a finite floating-point tensor of any shape.
mask has the same shape, dtype and device as x. Its values are 1 for keep and 0 for drop.
p is a scalar drop probability with 0 <= p < 1; the keep probability is 1 - p. The case p=1 is outside this exercise's contract.

Compute:

output = x * mask / (1 - p)

Return the same shape and dtype as x, leaving the inputs unchanged. Apply the supplied mask even if its observed fraction of zeros differs from p. Do not infer a new scaling factor from the mask or the number of kept entries.

For x=[-2, 4, 6], mask=[1, 0, 1] and p=0.5, return [-4, 0, 12]. For an all-zero mask, return zeros. For an all-one mask, still divide by 1 - p. With p=0, the expression reduces to x * mask.

Why scale the kept values?

If a random mask keeps each activation with probability 1 - p, then its average multiplier is (1 - p) / (1 - p) = 1. This preserves each activation's expected value over repeated random masks. It does not preserve the sum or mean for every individual mask.

During ordinary evaluation, inverted dropout is disabled and inputs pass through unchanged. This function has no evaluation-mode flag: it always applies the supplied training-time mask and scale.

Understand the ideas behind this exercise. The Problem tab contains the exact implementation requirements.

Choose initial weights that keep signals at a useful scale

Initialization helps keep activations and gradients from shrinking or growing rapidly across layers. Under suitable independence and variance assumptions, summing fan_in weighted inputs motivates weight variance proportional to 1 / fan_in.

Xavier initialization considers both input and output widths. He initialization also accounts for how ReLU removes the negative part of its input. Follow the specified formula and distribution rather than using these names as complete instructions.

A standard deviation is not the same as a uniform distribution's bound. Uniform samples from [-a, a] have variance a²/3. To get standard deviation s, the bound must be sqrt(3) * s.

Drop activations during training

With keep probability q, inverted dropout multiplies each activation x by m/q, where m is a zero-or-one mask sampled with probability q of being one. For q > 0, the expected output is x. A particular draw can still be very different. At evaluation, inverted dropout leaves the input unchanged.

For x = 4 and q = 0.5, the training output is either 0 or 8, each with probability one half. Their average is 4. The endpoint q = 0 needs explicit handling because the formula divides by zero.

Drop residual branches with stochastic depth

Stochastic depth drops an entire residual branch, often sharing a mask across its feature dimensions. The shortcut remains present. Some conventions scale surviving branches during training; others scale the branch at evaluation. Do not combine the two conventions.

Check mask shape, training/evaluation mode, and the required random seed or key behavior. Reusing a JAX key for the same random operation repeats its draw. PyTorch and JAX need not produce the same samples from the same integer seed.

Detecting runtime…

Implement Dropout

Your task

Why scale the kept values?

Choose initial weights that keep signals at a useful scale

Drop activations during training

Drop residual branches with stochastic depth

Your code, saved here

Stuck?

PyTorch reference solution