Skip to content

← JAX Stochasticity step 3 of 25

Easy Primitives

Bernoulli Mask Sampling

Jump to the task

Sample a yes-or-no decision at each position

A Bernoulli draw is True with probability p and False with probability 1 - p. For a mask, True means keep and False means remove. Here p is the keep probability, not the drop probability.

A 3-by-3 mask with p = 0.5 does not have a prescribed number of True entries. Its expected count is 9 * 0.5 = 4.5 over repeated draws; any single mask has an integer count that can differ from that average.

import jax
import jax.numpy as jnp

key = jax.random.PRNGKey(0)
mask_bool = jax.random.bernoulli(key, 0.5, (3, 3))
mask = mask_bool.astype(jnp.float32)
# [[0., 0., 1.],
#  [1., 0., 1.],
#  [1., 0., 0.]]

This example uses the default Threefry generator and sampling settings of the exercise runtime. It requests nine decisions in one array draw, not nine repeated calls with the same key. Repeating the same draw with the same key and settings reproduces the mask.

Why convert the mask?

Bernoulli returns a boolean array. This task explicitly asks for float32 zeros and ones, so use astype(jnp.float32). The conversion is an output requirement, not a prerequisite for all arithmetic: multiplying a floating-point array by a boolean mask already works in JAX.

Multiplying activations by this mask removes some entries. It is not the entire inverted-dropout operation: training-time inverted dropout also rescales kept values by 1 / p when p > 0. Do not add that scaling here; the requested output contains only zeros and ones.

The shape-array interface is for eager calls

This exercise supplies [H, W] as values in a JAX array. Convert the concrete values to Python integers before sampling. Directly wrapping the whole function with jax.jit would trace those values, so int(shape[0]) would fail. int() cannot make a runtime shape static, and a JAX array cannot simply be marked as a hashable static argument.

A JIT-oriented interface would accept a key and a Python tuple (H, W), with that tuple marked static. The sampling operation itself supports JIT when its output shape is known during tracing.

Problem

Implement bernoulli_mask(seed, p, shape) using one call to jax.random.bernoulli and return its mask as float32.

  • Seed is a whole Python number in [0, 2**32), supplied as a float. Create the key with jax.random.PRNGKey(int(seed)).
  • p is a floating-point scalar in [0, 1], the probability of a one.
  • shape is a length-two JAX array containing nonnegative whole values [H, W] as floats. Use (int(shape[0]), int(shape[1])) for the draw.
  • Return shape (H, W). p = 0 gives all zeros; p = 1 gives all ones. A zero dimension gives an empty array with the requested shape.

Preserve the runtime's default sampling configuration and use the original key without adding a split. Tests check the reproducible mask, not a target percentage of ones.