code: def make_causal_mask(
      x: Array, extra_batch_dims: int = 0, dtype: Dtype = jnp.float32
  ) -> Array:
      """Make a causal mask for self-attention.
  
      In case of 1d inputs (i.e., `[batch..., len]`, the self-attention weights
      will be `[batch..., heads, len, len]` and this function will produce a
      causal mask of shape `[batch..., 1, len, len]`.
  
      Args:
        x: input array of shape `[batch..., len]`
        extra_batch_dims: number of batch dims to add singleton axes for, none by
          default
        dtype: mask return dtype
  
      Returns:
        A `[batch..., 1, len, len]` shaped causal mask for 1d attention.
      """
      idxs = jnp.broadcast_to(
          jnp.arange(x.shape[-1], dtype=jnp.int32), x.shape
      )
      return make_attention_mask(
          idxs,
          idxs,
          jnp.greater_equal,
          extra_batch_dims=extra_batch_dims,
          dtype=dtype,
      )
  
imports: from __future__ import annotations
  import functools
  from typing import Any, Callable, Optional, Tuple, overload
  import jax
  import jax.numpy
  from jax import lax, random
  from flax.experimental import nnx
  from flax.experimental.nnx.nnx import rnglib
  from flax.experimental.nnx.nnx import flaglib
  from flax.experimental.nnx.nnx.module import Module, first_from
  from flax.experimental.nnx.nnx.nn import initializers
  from flax.experimental.nnx.nnx.nn.dtypes import promote_dtype
  from flax.experimental.nnx.nnx.nn.linear import DotGeneralT, LinearGeneral, PrecisionLike, default_kernel_init
  from flax.experimental.nnx.nnx.nn.normalization import LayerNorm
label_desc: Use image node_id 4 for calling a global function with example usage: make_causal_mask(x, extra_batch_dims, dtype) and returns: make_attention_mask