Set values in row to zero before index value of row [NumPy or Tensorflow]

A = np.array([0,1,1,0,1])
N = A.shape[0]
column = (A > 0).reshape((N, 1))
mask = np.ones((N, N), dtype=np.bool)
mask = np.where(column, False, np.tril(mask, -1))
mask = np.cumsum(mask, axis=0)
B = np.where(mask, 0, np.tile(A, (N, 1)))

[[0 1 1 0 1]
 [0 1 1 0 1]
 [0 1 1 0 1]
 [0 0 0 0 1]
 [0 0 0 0 1]]

Explanation

1. Compute lower triangular matrix

[[False False False False False]
 [ True False False False False]
 [ True  True False False False]
 [ True  True  True False False]
 [ True  True  True  True False]]

2. Find ones in A and fill corresponding rows with False

[[False False False False False]
 [False False False False False]
 [False False False False False]
 [ True  True  True False False]
 [False False False False False]]

3. Compute cumulative sum to set zeros to all rows below. This is a mask of all elements that should be zeroed out

[[0 0 0 0 0]
 [0 0 0 0 0]
 [0 0 0 0 0]
 [1 1 1 0 0]
 [1 1 1 0 0]]

4. Repeat array A N times

[[0 1 1 0 1]
 [0 1 1 0 1]
 [0 1 1 0 1]
 [0 1 1 0 1]
 [0 1 1 0 1]]

5. Mask its elements

[[0 1 1 0 1]
 [0 1 1 0 1]
 [0 1 1 0 1]
 [0 0 0 0 1]
 [0 0 0 0 1]]