Roman Novak

AFAIK @sschoenholz is working on an update to the batching code, which might address this, but he's on vacation now. If you have a simple reproducible example, I can also...

Sorry for delay/neurips! Could you please describe the dataset you're using, and the exact bottleneck you're hitting (too slow, or out of memory, ...)? I suspect potential improvements here might...

Thanks for clarifying and for the example, this is a great point, and it's easy to change the API with a simple change like this: (essentially just passing the `to_dense`...

FYI @sschoenholz and I have added a new argument `implementation` to `stax.Aggregate`, which supports the new `implementation="SPARSE"`, please see the commit message in https://github.com/google/neural-tangents/commit/b29337daf9a4e1f5b817f1689809021f52385f02 and the new docstring https://neural-tangents.readthedocs.io/en/latest/neural_tangents.stax.html#neural_tangents.stax.Aggregate, this...

Thank you for the detailed example and updates! To clarify, the issue goes away when using higher precision? This is a somewhat common issue, I don't have an immediate solution,...

Based on the numbers in the error message, am I correct that it's raised when running `python examples/infinite_fcn.py`? Unfortunately I can't reproduce it on my machine so far, with same...

Another avenue to investigate is to try running this python program ```python from scipy import linalg # if the above works, next try # from jax.scipy import linalg import numpy...

Thanks for the update! In this case I recommend submitting this issue to https://github.com/google/jax/issues with a repro from https://github.com/google/neural-tangents/issues/152#issuecomment-1121615513 since it appears to be a jaxlib rather than an NT...

Also, I suggest regularly syncing the fork (https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork) to make sure the PR stays in sync with latest changes.

Regarding tests, it looks like Github allows for 20 concurrent jobs running https://docs.github.com/en/actions/learn-github-actions/usage-limits-billing-and-administration#usage-limits, and we currently have 8 (x32/x64 * 4 python versions) linux jobs + 1 pytype job +...