docstr: 1's in the lower triangle, counting from the lower right corner. Same as tf.matrix_band_part(tf.ones([nd, ns]),
  -1, ns-nd), but doesn't produce garbage on TPUs.
docstr_tok: ["1","'s","in","the","lower","triangle",",","counting","from","the","lower","right","corner",".","Same","as","tf.matrix_band_part","(","tf.ones","(","[","nd",",","ns","]",")",",","-1",",","ns-nd",")",",","but","does","n't","produce","garbage","on","TPUs","."]
code: def causal_attention_mask(nd, ns, dtype):
      """
      1's in the lower triangle, counting from the lower right corner. Same as tf.matrix_band_part(tf.ones([nd, ns]),
      -1, ns-nd), but doesn't produce garbage on TPUs.
      """
      i = tf.range(nd)[:, None]
      j = tf.range(ns)
      m = i >= j - ns + nd
      return tf.cast(m, dtype)
  
imports: from dataclasses import dataclass
  from typing import List, Optional, Tuple, Union
  import numpy
  import tensorflow
  from ...activations_tf import get_tf_activation
  from ...modeling_tf_outputs import TFBaseModelOutputWithPastAndCrossAttentions, TFCausalLMOutputWithCrossAttentions, TFSequenceClassifierOutputWithPast
  from ...modeling_tf_utils import TFCausalLanguageModelingLoss, TFConv1D, TFModelInputType, TFPreTrainedModel, TFSequenceClassificationLoss, TFSequenceSummary, TFSharedEmbeddings, get_initializer, keras_serializable, unpack_inputs
  from ...tf_utils import shape_list, stable_softmax
  from ...utils import DUMMY_INPUTS, ModelOutput, add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward, logging, replace_return_docstrings
  from .configuration_gpt2 import GPT2Config
label_desc: Use image node_id 3 for calling the TFAttention obj's underlying member method code with example usage: obj.causal_attention_mask(nd, ns, dtype) and returns: tf