Closed tingwl0122 closed 3 months ago
This PR fixes the for-loop implementation for batched-input.
We don't even need vmap for Arnoldi since it actually did a low-rank approximation on inverse Hessian. Regular matrix multiplication should work.
vmap
Description
This PR fixes the for-loop implementation for batched-input.
1. Motivation and Context
We don't even need
vmap
for Arnoldi since it actually did a low-rank approximation on inverse Hessian. Regular matrix multiplication should work.2. Summary of the change
3. What tests have been added/updated for the change?