Can't run when using neural-tangents 0.3.2

[EdgeLLM]

When I use neural-tangents 0.3.2 to run the 1d_regression.ipynb, I got ValuError:

Could you help me with this?

[tancik]

Yep, it looks like the neural_tangents 0.3.2 has some breaking changes. I updated the notebook to use version 0.2.2.

[fengzlj]

Hello, When I use neural-tangents 0.2.2 to run the 1d_regression.ipynb, I got AttributeError: 'ShapedArray' object has no attribute 'val':

Could you help me with this?

[ShineChen1024]

Hello, When I use neural-tangents 0.2.2 to run the 1d_regression.ipynb, I got AttributeError: 'ShapedArray' object has no attribute 'val':

Could you help me with this?

you should try to use jax==0.1.76

[geyang]

@chenshine1 when I use jax==0.1.76 I get the following error

<ipython-input-22-f42eedd5376b> in <module>()
----> 1 import jax
      2 from jax import random, grad, jit, vmap
      3 from jax.config import config
      4 from jax.lib import xla_bridge
      5 import jax.numpy as np

4 frames
/usr/local/lib/python3.7/dist-packages/jax/lib/__init__.py in <module>()
     58   pytree = xla_client._xla.pytree
     59   jax_jit = xla_client._xla.jax_jit
---> 60 from jaxlib import cusolver
     61 try:
     62   from jaxlib import cuda_prng

ImportError: cannot import name 'cusolver' from 'jaxlib' (/usr/local/lib/python3.7/dist-packages/jaxlib/__init__.py)