noahgolmant / pytorch-hessian-eigenthings

Efficient PyTorch Hessian eigendecomposition tools!
MIT License
360 stars 43 forks source link

Stochastic lanczos quadrature #36

Open devansh20la opened 3 years ago

devansh20la commented 3 years ago

Changes

  1. Added stochastic_lanczos.py under hessian_eigenthings. As per the discussions on the OG repo, the hvp operations are performed on gpu while the rest of the computations are performed on the CPU with float64 precision.
  2. tests/spectral_density_tests.py. As of now I only perform statistical distance test here. I wanted to make this pull request and get the conversation going.
  3. General housekeeping in "__init__.py"