Skip to content

← Flax step 20 of 100

Medium Primitives

Implement Dropout with RNG Threading

Jump to the task

Implement dropout directly inside a Flax Linen module. You will obtain a random key with self.make_rng, sample a mask, and scale the retained values. This module has no learned parameters and does not save the mask as state.

Keep some values, scale the survivors

Let rate be the probability of dropping an element and keep_prob = 1 - rate. For 0 < rate < 1, sample one boolean per input element, independently:

mask = jax.random.bernoulli(rng, p=keep_prob, shape=x.shape)
output = jnp.where(mask, x / keep_prob, 0.0)

A true mask entry keeps the corresponding value and scales it up. For example, with rate=0.5, an input value of 3 becomes either 6 or 0. Its expected value is 0.5 * 6 + 0.5 * 0 = 3. This scaling is called inverted dropout. It preserves each element's expected value, not the sum of every sampled output. A short vector can have all its elements dropped or all retained.

Evaluation uses deterministic=True, which returns x unchanged. Rate zero also returns x. With dropout enabled, rate one returns all zeros: handle this case before dividing, because its keep probability is zero.

Obtain a key from a named stream

Inside the module, request a key with:

rng = self.make_rng("dropout")

The caller supplies the stream through rngs={"dropout": drop_key} on apply. Linen derives the requested key from the stream key, module path and call counter. It is not simply the top-level drop_key passed through unchanged. Keep the module structure fixed when reproducing a seeded mask.

If a requested stream is absent, Linen can fall back to an available params stream. If neither is supplied, requesting the key fails. This exercise uses an explicit dropout stream to make the source of randomness clear.

make_rng needs a bound module with access to RNGs, not specifically a method named __call__. It can also be used in another method invoked through apply. For this exercise, put the calculation in a compact __call__.

Repeated calls and initialization

The same key, input shape and call structure reproduce the same mask. This is useful for testing. For fresh randomness across training steps, split a carried key and pass a new dropout key to each call. Linen's internal key derivation does not automatically advance the caller's key between separate apply calls.

Initializing this parameter-free module returns an empty variables mapping. The params key in the exercise's initialization recipe is an RNG stream, not evidence that any parameter arrays exist. There is no dropout variable collection to register. The identity and all-zero branches need no RNG.

Problem

Implement MyDropout(nn.Module) with a rate field and a compact __call__(self, x, deterministic):

  1. Return x if deterministic is true or rate == 0.0.
  2. Otherwise, return jnp.zeros_like(x) if rate == 1.0.
  3. Otherwise, obtain one key with self.make_rng("dropout"), sample the Bernoulli mask shown above, and return the scaled, masked output.

Implement my_dropout_forward(seed, x, rate, deterministic) around it:

  • Construct MyDropout(rate=float(rate)) and convert the flag with bool(deterministic >= 0.5).
  • Set init_key = jax.random.PRNGKey(int(seed)) and drop_key = jax.random.PRNGKey(int(seed) + 1) without splitting them.
  • Initialize with {"params": init_key, "dropout": drop_key} and deterministic=False.
  • Apply with the returned variables mapping, the converted flag and rngs={"dropout": drop_key}.

x is a finite float32 JAX vector of length at least one. seed is a nonnegative integer-valued scalar that may arrive as a float. rate is in [0, 1] and deterministic is 0.0 or 1.0. Return a JAX vector with the same shape as x, leaving x unchanged. Do not wrap another Dropout module: that would change the module path used to derive the random key.