Package: matfree
matfree: Matrix-free linear algebra in JAXRandomised and deterministic matrix-free methods for trace estimation, matrix functions, and/or matrix factorisations.
Builds onJAX.⚡ Stochastictrace estimationincluding batching, control variates, and uncertainty quantification⚡ A stand-alone implementation ofstochastic Lanczos quadrature⚡ Matrix-decomposition algorithms forlarge sparse eigenvalue problems⚡ Polynomial methods for approximatingfunctions of large matricesand many other things.
Everything is natively compatible with the rest of JAX:
JIT compilation, automatic differentiation, vectorisation, and PyTrees.Let us know what you think about matfree!InstallationTo install the package, runpip install matfreeImportant:This assumes you already have a working installation of JAX.
To install JAX,follow these instructions.
To combinematfreewith a CPU version of JAX, runpip install matfree[cpu]which is equivalent to combiningpip install jax[cpu]withpip install matfree.
(But do not only use matfree on CPU!)Minimal exampleImport matfree and JAX, and set up a test problem.>>>importjax>>>importjax.numpyasjnp>>>frommatfreeimporthutchinson>>>>>>jnp.set_printoptions(1)>>>A=jnp.reshape(jnp.arange(12.0),(6,2))>>>>>>defmatvec(x):...returnA.T@(A@x)...Estimate the trace of the matrix:>>># Determine the shape of the base-samples>>>input_like=jnp.zeros((2,),dtype=float)>>>sampler=hutchinson.sampler_rademacher(input_like,num=10_000)>>>>>># Set Hutchinson's method up to compute the traces>>># (instead of, e.g., diagonals)>>>integrand=hutchinson.integrand_trace(matvec)>>>>>># Compute an estimator>>>estimate=hutchinson.hutchinson(integrand,sampler)>>># Estimate>>>key=jax.random.PRNGKey(1)>>>trace=jax.jit(estimate)(key)>>>>>>print(trace)508.9>>>>>># for comparison:>>>print((jnp.trace(A.T@A)))506.0TutorialsFind many more tutorials inMatfree's documentation.These tutorials include, among other things:Log-determinants:Use stochastic Lanczos quadrature to compute matrix functions.Pytree-valued states:Combining neural-network Jacobians with stochastic Lanczos quadrature.Control variates:Use control variates and multilevel schemes to reduce variances.Higher moments and UQ:Compute means, variances, and other moments simultaneously.Vector calculus:Use matrix-free linear algebra to implement vector calculus.Let us knowwhat you use matfree for!Use Matfree's continuous integrationTo install all test-related dependencies, (assuming JAX is installed; if not, runpip install .[cpu]), executepip install .[test]Then, run the tests viamake testInstall all formatting-related dependencies viapip install .[format]and format the code viamake formatTo lint the code, install the pre-commit hookpip install .[lint]
pre-commit installand run the linters viamake lintInstall the documentation-related dependencies aspip install .[doc]Preview the documentation viamake doc-previewand check whether the docs build correctly viamake doc-buildContribute to MatfreeContributions are absolutely welcome!Issues:Most contributions start with an issue.
Please don't hesitate tocreate issuesin which you
ask for features, give feedback on performances, or simply want to reach out.Pull requests:To make a pull request, proceed as follows:Fork the repository.Install all dependencies withpip install .[full]orpip install -e .[full].Make your changes.From the root of the project, run the tests viamake test, and check outmake formatandmake lintas well. Use the pre-commit hook if you like.When making a pull request, keep in mind the following (rough) guidelines:Most PRs resolve an issue.Most PRs contain a single commit.Here is how we can write better commit messages.Most enhancements (e.g. new features) are covered by tests.Extend Matfree's documentationWrite a new tutorial:To add a new tutorial, create a Python file intutorials/and fill it with content.
Use docstrings (mirror the style in the existing tutorials).
Make sure to satisfy the formatter and linter.
That's all.Then, the documentation pipeline will automatically convert those into a format compatible
with Jupytext, which subsequently includes it into the documentation.
If you do not want to make the tutorial part of the documentation, make the filename
have a leading underscore.Extend the developer documentation:To extend the developer documentation, create a new section in the README.
Use a second-level header (a header that starts with "##") and fill the section
with content.
Then, the documentation pipeline will turn this section into a page in the developer documentation.Create a new module:To make a new module appear in the documentation, create the new module inmatfree/,
and fill it with content.
Unless the module name starts with an underscore or is placed in the backend,
the documentation pipeline will take care of the rest.