code: def dropout_impl(
      rng: PRNGKey,
      rate: float,
      x: jax.Array,
      broadcast_dims: Sequence[int] = (),
  ) -> jax.Array:
      """See dropout."""
      try:
          if rate < 0 or rate >= 1:
              raise ValueError("rate must be in [0, 1).")
  
          if rate == 0.0:
              return x
      except jax.errors.ConcretizationTypeError:
          pass
  
      broadcast_shape = list(x.shape)
      for dim in broadcast_dims:
          if dim > len(broadcast_shape):
              raise ValueError(
                  "Broadcast dimension does not exist. Got dimension "
                  f"{dim} for shape {broadcast_shape}."
              )
          broadcast_shape[dim] = 1
  
      keep_rate = 1.0 - rate
      keep = jax.random.bernoulli(rng, keep_rate, shape=broadcast_shape)
      keep = jnp.broadcast_to(keep, x.shape)
      return keep * x / keep_rate
  
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 9 for calling a global function with example usage: dropout_impl(rng, rate, x, broadcast_dims) and returns: unknown, x