Triangle matrix operations with 4D matrix

[markub3327]

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?

[markub3327]

Solution:

# (Timestep, Patch)
def create_look_ahead_mask(size_0, size_1):
  n = int(size_0 * (size_0+1) / 2)
  mask = tfp.math.fill_triangular(tf.ones((n,), dtype=tf.int32), upper=False)
  mask = mask[:, tf.newaxis, :, tf.newaxis]  # (timestep, 1, timestep, 1)
  return mask