@staticmethod
def causal_attention_mask(nd, ns, dtype):
    "1's in the lower triangle, counting from the lower right corner.\n        Same as tf.matrix_band_part(tf.ones([nd, ns]), -1, ns-nd), but doesn't produce garbage on TPUs.\n        "
    i = tf.range(nd)[:, None]
    j = tf.range(ns)
    m = (i >= ((j - ns) + nd))
    return tf.cast(m, 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.
@staticmethod
def causal_attention_mask(nd, ns, dtype):
    "1's in the lower triangle, counting from the lower right corner.\n        Same as tf.matrix_band_part(tf.ones([nd, ns]), -1, ns-nd), but doesn't produce garbage on TPUs.\n        "
    i = tf.range(nd)[:, None]
    j = tf.range(ns)
    m = (i >= ((j - ns) + nd))
    return tf.cast(m, dtype)