class_docstr: Lazy causal mask, prevents the model from attending to future tokens.
  
  Attributes:
    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.
class_docstr_tok: ["Lazy","causal","mask",",","prevents","the","model","from","attending","to","future","tokens",".","Attributes",":","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","."]
code: def __init__(
      self,
      shape: Tuple[int, int],
      offset: int = 0,
      shard_count: int = 1,
  ):
      self.offset = offset
  
      def causal_mask_function(q_ids, kv_ids):
          # When evaluating the mask in _process_mask we typically work with numpy
          # array views.
          # Avoid the addition when possible to avoid instantiating an actual array.
          if self.offset == 0:
              return q_ids >= kv_ids
          else:
              return q_ids + self.offset >= kv_ids
  
      mask_function = causal_mask_function
  
      super().__init__(
          shape=shape,
          mask_function=mask_function,
          shard_count=shard_count,
      )
  
imports: from __future__ import annotations
  import dataclasses
  from typing import Any, Callable, Sequence, Tuple
  import numpy
label_desc: Use image node_id 1 to create a new CausalMask object from inherited base classes: _ComputableMask with example: obj = CausalMask(shape, offset, shard_count)