Easy End-to-End

Weight Initialization

Compute two standard deviations used to initialize neural-network weights. This exercise returns the scales only: do not generate random weights or create a weight matrix.

Your task

Implement compute_init_scales(fan_in, fan_out) for positive integer inputs. For a dense layer, fan_in is the number of input features and fan_out is the number of output features. These counts are supplied directly, so you do not need to infer them from a matrix's storage layout.

Use these exact conventions:

xavier_scale = sqrt(2 / (fan_in + fan_out))
he_scale = sqrt(2 / fan_in)

The first is the standard deviation for Xavier normal initialization with gain 1. The second is the standard deviation for He normal initialization using fan-in scaling for ReLU. Do not add an activation-specific gain or substitute fan-out scaling.

Return a dictionary with keys "xavier_scale" and "he_scale", each holding a Python float rounded to seven decimal places. Use ordinary division before taking the square root; do not use integer division.

Example

For fan_in = 2 and fan_out = 2, the Xavier scale is sqrt(1/2) and the He scale is sqrt(1). Return:

{"xavier_scale": 0.7071068, "he_scale": 1.0}

Xavier depends on both counts. Swapping fan_in and fan_out leaves it unchanged, but generally changes the He scale.

What does scale mean?

For these normal initializers, scale means standard deviation, not variance or a hard bound on sampled weights. A zero-mean normal draw with standard deviation s has variance s**2 and can lie outside [-s, s]. Uniform initializer bounds use different formulas.

These are specified conventions for this exercise, not a guarantee that one formula is best for every model or activation. See the initialization reference for the gain and fan-mode options used by the library functions.

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…