keras/src/backend/tensorflow/numpy.py/triu
def triu(x, k=0):
    x = convert_to_tensor(x)

    def _positive_k_branch():
        shape = tf.shape(x)
        rows, cols = shape[-2], shape[-1]
        i, j = tf.meshgrid(tf.range(rows), tf.range(cols), indexing="ij")
        mask = i <= j - k
        return tf.where(tf.broadcast_to(mask, shape), x, tf.zeros_like(x))

    return tf.cond(
        k <= 0, lambda: tf.linalg.band_part(x, -k, -1), _positive_k_branch
    )