Given the tensor A with shape (b,n,d) and another matrix M with shape (b,n), what is the most efficient way to split A into two tensors B and C with shapes (b,l,d) and (b,m,d), respecitvely, where n=l+m and the row A[k,i,:] should be included in tensor B if M[k,i] == 1, otherwise it should be included  in C.
The trivial solution is to have a for loop and selects the row in A based on the corresponding values in M. Is there any way to do this splitting in parallel?