As we all known, until now, GPU version eigs is not supported in JAX. In this repo is a simple implementation of eigs in JAX by rescuing useful code in TensorNetwork by google.
Copy jaxeigs.py to your project.
See test.py for usage.
avalokitesvara/jaxeigs
Folders and files
| Name | Name | Last commit date | ||
|---|---|---|---|---|