code: def _attention_dropout(
      attn_weights: jnp.ndarray,
      *,
      rate: float,
      broadcast: bool = True,
      dropout_rng: jnp.ndarray,
  ) -> jnp.ndarray:
      """Applies dropout on attention weights.
  
      This always applies the dropout. There is no `deterministic` parameter.
  
      Arguments:
        attn_weights: Attention weights.
        rate: The dropout rate. (_not_ the keep rate!)
        broadcast: Whether to broadcast on first and second last axis.
        dropout_rng: RNG.
  
      Returns:
        Weights after dropout.
      """
      keep_prob = jax.lax.tie_in(attn_weights, 1.0 - rate)
      if broadcast:
          # Dropout is broadcast across the batch+head+non-attention dimension.
          dropout_shape = list(attn_weights.shape)
          dropout_shape[0] = 1  # Broadcast batch.
          dropout_shape[-2] = 1  # Broadcast heads.
          keep = jax.random.bernoulli(
              dropout_rng, keep_prob, dropout_shape
          )
      else:
          keep = jax.random.bernoulli(
              dropout_rng, keep_prob, attn_weights.shape
          )
      multiplier = keep.astype(attn_weights.dtype) / jnp.asarray(
          keep_prob, dtype=attn_weights.dtype
      )
      return attn_weights * multiplier
  
imports: import functools
  from typing import Callable, Optional, Sequence, Tuple
  import flax.linen
  import jax
  import jax.numpy
  import numpy
  from scenic.model_lib.layers import nn_layers
label_desc: Use image node_id 1 for calling a global function with example usage: _attention_dropout(attn_weights) and returns: unknown