code: def make_causal_mask(x, dtype=jnp.float32):
      idxs = jnp.broadcast_to(
          jnp.arange(x.shape[-1], dtype=jnp.int32), x.shape
      )
      return make_attention_mask(
          idxs, idxs, jnp.greater_equal, dtype=dtype
      )
  
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 3 for calling a global function with example usage: make_causal_mask(x, dtype) and returns: make_attention_mask