code: def make_causal_mask(
      shape: Tuple[int, int], offset: int = 0
  ) -> np.ndarray:
      """Makes a causal attention mask.
  
      Args:
        shape: Shape of the 2-dim mask: (q_seq_len, kv_seq_len).
        offset: Offset of q start wrt kv. A positive offset shifts the bottom
          triangle upward, a negative one shifts it downward. A negative offset
          makes the first 'offset' rows of the attention matrix all 0s which leads
          to undefined softmax.
  
      Returns:
        The causal mask.
      """
      q_seq_len, kv_seq_len = shape
      q_idx = np.arange(q_seq_len, dtype=np.int32)
      kv_idx = np.arange(kv_seq_len, dtype=np.int32)
      return (q_idx[:, None] + offset >= kv_idx[None, :]).astype(
          np.bool_
      )
  
imports: from __future__ import annotations
  import dataclasses
  from typing import Any, Callable, Sequence, Tuple
  import numpy
label_desc: Use image node_id 1 for calling a global function with example usage: make_causal_mask(shape, offset) and returns: unknown