CAVEAT: Requires defined shape and can"t work with unknown shape.
size = np.prod(y.shape.as_list())
N = int(np.sqrt(size / self.num_matrices))
reshaped = tf.reshape(y, shape=(self.num_matrices, N, N))
indices = np.array([np.hstack(x) for x in
itertools.product(np.arange(self.num_matrices), np.dstack(np.tril_indices(N))[0])])
triangular = tf.reshape(tf.gather_nd(reshaped, indices), shape=[-1])
return triangular
After Change
if self.squeeze:
y = tf.expand_dims(y, axis=0)
indices = np.vstack(np.tril_indices(self.N)).T
indices = itertools.product(np.arange(self.num_matrices), indices)
indices = np.array([np.hstack(x) for x in indices])
triangular = tf.gather_nd(y, indices)
return tf.reshape(triangular, [self.num_matrices, (self.N**2 + self.N) // 2])
