class_docstr: Lazy local mask, prevents model from attending to tokens outside window.
  
  Attributes:
    _shape: Shape of the 2-dim mask: (q_seq_len, kv_seq_len).
    window_size: Size of the two sides of the local window (None identifes no
      limit for the given side).
    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.
    _q_sequence: Important for performance.
class_docstr_tok: ["Lazy","local","mask",",","prevents","model","from","attending","to","tokens","outside","window",".","Attributes",":","_shape",":","Shape","of","the","2-dim","mask",":","(","q_seq_len",",","kv_seq_len",")",".","window_size",":","Size","of","the","two","sides","of","the","local","window","(","None","identifes","no","limit","for","the","given","side",")",".","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",".","_q_sequence",":","Important","for","performance","."]
code: def __getitem__(self, idx) -> np.ndarray:
      if len(idx) != 2:
          raise NotImplementedError(f"Unsupported slice: {idx}")
      q_slice, kv_slice = idx
      if not isinstance(q_slice, slice) or not isinstance(
          kv_slice, slice
      ):
          raise NotImplementedError(f"Unsupported slice: {idx}")
  
      q_slice = _fill_slice(q_slice, self.shape[0])
      kv_slice = _fill_slice(kv_slice, self.shape[1])
  
      if self._q_sequence is None:
          rows = np.arange(q_slice.start, q_slice.stop)
      else:
          rows = self._q_sequence[q_slice]
  
      cols = np.arange(kv_slice.start, kv_slice.stop)
  
      left_size, right_size = self.window_size
  
      if left_size is None and right_size is None:
          return np.ones((rows.shape[0], cols.shape[0]), dtype=np.bool_)
      else:
          expanded_cols = cols[None, :]
          if self.offset != 0:
              expanded_rows = rows[:, None] + self.offset
          else:
              expanded_rows = rows[:, None]
          if left_size is not None and right_size is not None:
              return (expanded_rows <= expanded_cols + left_size) & (
                  expanded_cols - right_size <= expanded_rows
              )
  
          elif left_size is not None and right_size is None:
              return expanded_rows <= expanded_cols + left_size
          else:
              assert left_size is None and right_size is not None
              return expanded_cols - right_size <= expanded_rows
  
imports: from __future__ import annotations
  import dataclasses
  from typing import Any, Callable, Sequence, Tuple
  import numpy
label_desc: Use image node_id 3 for calling the LocalMask obj's underlying member method code with example usage: obj.__getitem__(idx) and returns: np, unknown, unknown, unknown