Issue with jax versions

[ritvikmahajan01]

I have jax 0.5.2 and jaxlib 0.5.1 (the default versions that are installed using pip). But when trying to import the library:

from neural_tangents import stax

I get the error:

AttributeError: module 'jax.random' has no attribute 'KeyArray'

It looks like this had been depreciated in newer jax version so I installed jax==0.4.23. However, then I am unable to install the corresponding jaxlib version:

ERROR: Could not find a version that satisfies the requirement jaxlib==0.4.23 (from versions: 0.4.34, 0.4.35, 0.4.36, 0.4.38, 0.5.0, 0.5.1)

How can I fix this?

[soldasim]

I am having the same issue. When I try to evaluate from neural_tangents import stax, I get

>>> from neural_tangents import stax
Traceback (most recent call last):
  File "<python-input-1>", line 1, in <module>
    from neural_tangents import stax
  File "/home/soldasim/.conda/envs/ibnn2/lib/python3.13/site-packages/neural_tangents/__init__.py", line 22, in <module>
    from . import predict
  File "/home/soldasim/.conda/envs/ibnn2/lib/python3.13/site-packages/neural_tangents/predict.py", line 17, in <module>
    from ._src.predict import (
    ...<9 lines>...
    )
  File "/home/soldasim/.conda/envs/ibnn2/lib/python3.13/site-packages/neural_tangents/_src/predict.py", line 43, in <module>
    from .utils import dataclasses, utils
  File "/home/soldasim/.conda/envs/ibnn2/lib/python3.13/site-packages/neural_tangents/_src/utils/utils.py", line 26, in <module>
    from .typing import Axes, PyTree
  File "/home/soldasim/.conda/envs/ibnn2/lib/python3.13/site-packages/neural_tangents/_src/utils/typing.py", line 67, in <module>
    class InitFn(Protocol):
    ...<13 lines>...
        ...
  File "/home/soldasim/.conda/envs/ibnn2/lib/python3.13/site-packages/neural_tangents/_src/utils/typing.py", line 77, in InitFn
    rng: random.KeyArray,
         ^^^^^^^^^^^^^^^
  File "/home/soldasim/.conda/envs/ibnn2/lib/python3.13/site-packages/jax/_src/deprecations.py", line 55, in getattr
    raise AttributeError(f"module {module!r} has no attribute {name!r}")
AttributeError: module 'jax.random' has no attribute 'KeyArray'
>>> 

I was able to downgrade to jax==0.4.23 jaxlib==0.4.23 but then I just get a different error.