Build one dropout keep mask per layer from a single integer seed, following an exact key-splitting rule.
Dropout randomly zeroes activations during training. A keep mask holds True for every entry that survives, and each entry is kept with probability keep_prob. The masks are fully determined by the seed and by the way keys are split, so follow this rule exactly.
Return a bool JAX array of shape (num_layers, followed by the entries of shape), for example (3, 2, 4) for three layers of shape [2, 4], with the mask of layer i at index i of the first axis.
Input: seed = 0, num_layers = 2, shape = [2,4], keep_prob = 0.5
Output: [[[true,true,false,true],[true,true,true,false]],[[true,true,true,false],[true,false,false,true]]]
Explanation: The seed 0 key is split once for layer 0, and the first key from that split is split again for layer 1. The two layer keys are different, so the two masks differ.
Input: seed = 42, num_layers = 3, shape = [5], keep_prob = 0.8
Output: [[true,true,true,true,true],[true,true,true,true,false],[true,true,true,true,false]]
Input: seed = 2026, num_layers = 1, shape = [3,3], keep_prob = 0.3
Output: [[[true,true,false],[false,false,false],[false,true,true]]]
key, subkey = jax.random.split(key) gives a fresh subkey and a new running key.
jax.random.bernoulli(subkey, p, shape) returns a bool array of that shape.
jnp.stack(list_of_masks) adds a new first axis.
Accepts: integer
Accepts: integer
Accepts: array
Accepts: number
Build one dropout keep mask per layer from a single integer seed, following an exact key-splitting rule.
Dropout randomly zeroes activations during training. A keep mask holds True for every entry that survives, and each entry is kept with probability keep_prob. The masks are fully determined by the seed and by the way keys are split, so follow this rule exactly.
Return a bool JAX array of shape (num_layers, followed by the entries of shape), for example (3, 2, 4) for three layers of shape [2, 4], with the mask of layer i at index i of the first axis.
Input: seed = 0, num_layers = 2, shape = [2,4], keep_prob = 0.5
Output: [[[true,true,false,true],[true,true,true,false]],[[true,true,true,false],[true,false,false,true]]]
Explanation: The seed 0 key is split once for layer 0, and the first key from that split is split again for layer 1. The two layer keys are different, so the two masks differ.
Input: seed = 42, num_layers = 3, shape = [5], keep_prob = 0.8
Output: [[true,true,true,true,true],[true,true,true,true,false],[true,true,true,true,false]]
Input: seed = 2026, num_layers = 1, shape = [3,3], keep_prob = 0.3
Output: [[[true,true,false],[false,false,false],[false,true,true]]]
key, subkey = jax.random.split(key) gives a fresh subkey and a new running key.
jax.random.bernoulli(subkey, p, shape) returns a bool array of that shape.
jnp.stack(list_of_masks) adds a new first axis.
Accepts: integer
Accepts: integer
Accepts: array
Accepts: number