Recently I met a problem with TE. As we all know, TE uses a declarative way to define a computation node, where lambda i, j: xxx means output[i, j] = xxx.
I wonder how to get output according to a mask using TE, such as output[mask[i], j] = xxx and the other elements are 0.