code: def dropout(cfg: Config, x, broadcast_dims=(-2,), *, rngs: nnx.Rngs):
      if cfg.dropout_rate == 0.0:
          return x
      broadcast_shape = list(x.shape)
      for dim in broadcast_dims:
          broadcast_shape[dim] = 1
      keep_rate = 1.0 - cfg.dropout_rate
      key = rngs.dropout()
      mask = jax.random.bernoulli(
          key, p=keep_rate, shape=broadcast_shape
      )
      return jax.lax.select(
          jnp.broadcast_to(mask, x.shape),
          x / keep_rate,
          jnp.zeros_like(x),
      )
  
imports: import dataclasses
  import typing
  import jax
  import jax.numpy
  import numpy
  from jax.sharding import PartitionSpec
  from flax.experimental import nnx
label_desc: Use image node_id 8 for calling a global function with example usage: dropout(cfg, x, broadcast_dims) and returns: jax, x