How to implement 'Matrix exponential times a vector'?

1

The documentation here explains how to regularize the weights of different layers. Specifically, there is an example of matrix_exp operation over the weights of a layer. I am interested in implementing a matrix exponential times a vector, that is, exp(tW)z where t is a batch of different time steps (shape: b x 1) and W are the parameters of the layer (say a linear layer with d x d), and z is a latent representation of a batch, of a shape: b x d.
Using the link above, I can implement exp(tW) for a specific t for the whole batch (1x1) via building a new module and then feeding it with my batch of latent representations z. Is it possible to do all in one operation?

In particular, I am interested in this implementation; Computing the Action of the Matrix Exponential, with an Application to Exponential Integrators by Al-Mohy et al.

2

You might consider seeing if your module (even if implemented as multiple operations) can be fused into a single operation via e.g., torch.compile: torch.compile Tutorial — PyTorch Tutorials 2.0.1+cu117 documentation.