The mask in reduce_to_vector is a Vector, while the input is a Matrix, making it impossible to apply the mask prior to the operation. If select_by_mask could do vec->matrix broadcasting, the mask could be applied earlier and possibly save compute time.
Alternatively, the mask could be brought into the linalg.generic loop and used to skip entire rows.
The 2nd approach would be faster and simpler if it works.
The mask in
reduce_to_vectoris a Vector, while the input is a Matrix, making it impossible to apply the mask prior to the operation. Ifselect_by_maskcould do vec->matrix broadcasting, the mask could be applied earlier and possibly save compute time.Alternatively, the mask could be brought into the
linalg.genericloop and used to skip entire rows.The 2nd approach would be faster and simpler if it works.