def __call__(
      self, value: named_axes.NamedArray, /, **side_inputs
  ) -> named_axes.NamedArray:
    """Randomly drops out components of the input."""
    rng = side_inputs[self.random_stream_input_name]
    random_mask_axes = {
        name: size
        for name, size in value.named_shape.items()
        if name not in self.share_across_axes
    }
    drop_mask = named_axes.wrap(
        jax.random.bernoulli(
            rng.next_key(),
            p=self.drop_rate,
            shape=tuple(random_mask_axes.values()),
        ),
        *random_mask_axes.keys(),
    )
    return named_axes.nmap(jnp.where)(
        drop_mask, 0, value / (1 - self.drop_rate)
    )
Randomly drops out components of the input.
penzai/nn/dropout.py