code: def dropout(
      rng: PRNGKey,
      rate: float,
      x: jax.Array,
      broadcast_dims: Sequence[int] = (),
  ) -> jax.Array:
      """Randomly drop units in the input at a given rate.
  
      See: http://www.cs.toronto.edu/~hinton/absps/dropout.pdf
  
      Args:
        rng: A JAX random key.
        rate: Probability that each element of ``x`` is discarded. Must be a scalar
          in the range ``[0, 1)``.
        x: The value to be dropped out.
        broadcast_dims: specifies dimensions that will share the same dropout mask.
  
      Returns:
        x, but dropped out and scaled by ``1 / (1 - rate)``.
  
      Note:
        This involves generating `x.size` pseudo-random samples from U([0, 1))
        computed with the full precision required to compare them with `rate`. When
        `rate` is a Python float, this is typically 32 bits, which is often more
        than what applications require. A work-around is to pass `rate` with a lower
        precision, e.g. using `np.float16(rate)`.
      """
      return dropout_impl(rng, rate, x, broadcast_dims=broadcast_dims)
  
imports: from collections.abc import Iterable, Sequence
  import functools
  from typing import Any, Callable, Optional
  from haiku._src import base
  from haiku._src import initializers
  from haiku._src import module
  from haiku._src import typing
  from haiku._src.typing import PRNGKey
  import jax
  from jax import lax
  import jax.numpy
  import numpy
label_desc: Use image node_id 8 for calling a global function with example usage: dropout(rng, rate, x, broadcast_dims) and returns: dropout_impl