I’m trying to do some optimizations. My understanding is that, usually, it’s better to vectorize than to do loops. But this doesn’t seem to be the case for matrix_exp. Here are some tests:
matrices = [torch.rand(1000, 1000) for _ in range(16)]

Test using loops:

%%timeit
for m in matrices:
    e = torch.matrix_exp(m)
2.24 s ± 48.4 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Test with vectorization:

%%timeit
torch.matrix_exp(torch.stack(matrices))
2.21 s ± 16.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

I would love to understand why there’s very little speed-up in this case, and if there’s anything else I can do.