I am modifying identity_3d which is initialized as an n by n by n numpy array per the following operations:

     identity_3d = np.zeros((n, n, n))


        idx = np.arange(n)
        identity_3d[:, idx, idx] = 1 
        I,J = np.nonzero(wR==0)
        identity_3d[I,:,J]=0
        identity_3d[I,J,:]=0

If identity_3d was an  Tensor instead, is there a way to perform the equivalent operation?