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.