Select random non-zero element from each row of 3d matrix in non-eager tensorflow

Question

I am trying to implement the triplet loss with random negative triplet selection. Right now I have a tensor of shape (batch_size, batch_size, batch_size) where element (i,j,