Hi,
I have problem that requires applying function objects (like torch.sin) to some columns of a matrix X. Right now I solve this by having a list of tuples like this: [(function, col)] where function is the function object and col is the column this function acts on.
In each step I stack the resulting columns and return the matrix. Now I’m not sure this is the most performant way because it’s doing a list comprehension over all columns for every evaluation. Doing something like
X[:,i] = torch.sin(X[:,i])
is also not possible because it breaks autograd. I was thinking the new torch.vmap could somehow help here?
Another problem with my approach is that due to the list mentioned above I can’t jit the module because the functor datatype is not supported.
Do you have any tips on how to improve the implementation?