Triangle matrix operations with 4D matrix

Hello,

this code:

def create_look_ahead_mask(size_0):
  mask = np.ones((size_0, size_0), dtype=np.int32)

  for a in range(size_0):          # Timestep
      for c in range(size_0):      # Timestep
          if c > a:
            mask[a, c] = 0

  return mask  # (seq0_len, seq1_len, seq0_len, seq1_len)

print(create_look_ahead_mask(3), "\n")

is equivalent to:

def create_look_ahead_mask(size):
  n = int(size * (size+1) / 2)
  mask = tfp.math.fill_triangular(tf.ones((n,), dtype=tf.int32), upper=False)
  return mask

print(create_look_ahead_mask(3))

Here is another code very close to the above:

def create_look_ahead_mask(size_0, size_1):
  mask = np.ones((size_0, size_1, size_0, size_1), dtype=np.int32)

  for a in range(size_0):          # Timestep
    for b in range(size_1):        # Patch
      for c in range(size_0):      # Timestep
        for d in range(size_1):    # Patch
          if c > a:
            mask[a, b, c, d] = 0

  return mask  # (seq0_len, seq1_len, seq0_len, seq1_len)

print(create_look_ahead_mask(3, 3), "\n")

But what is the equivalent to it with using tfp.math.fill_triangular or another matrix operation?

No repository file or test is named. Start by comparing the NumPy loop implementation with tfp.math.fill_triangular and other TensorFlow matrix operations, then verify that the chosen operation reproduces the requested (size_0, size_1, size_0, size_1) mask for the provided examples.