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

[Zethson]

Report

Matplotlib is building the font cache; this may take a moment.
Traceback (most recent call last):
  File "<string>", line 1, in <module>
  File "/home/docs/checkouts/readthedocs.org/user_builds/pertpy/checkouts/529/pertpy/__init__.py", line 18, in <module>
    from . import plot as pl
  File "/home/docs/checkouts/readthedocs.org/user_builds/pertpy/checkouts/529/pertpy/plot/__init__.py", line 2, in <module>
    from pertpy.plot._coda import CodaPlot as coda
  File "/home/docs/checkouts/readthedocs.org/user_builds/pertpy/checkouts/529/pertpy/plot/_coda.py", line 13, in <module>
    from pertpy.tools._coda._base_coda import CompositionalModel2
  File "/home/docs/checkouts/readthedocs.org/user_builds/pertpy/checkouts/529/pertpy/tools/__init__.py", line 2, in <module>
    from pertpy.tools._cinemaot import Cinemaot
  File "/home/docs/checkouts/readthedocs.org/user_builds/pertpy/checkouts/529/pertpy/tools/_cinemaot.py", line 12, in <module>
    from ott.geometry import pointcloud
  File "/home/docs/checkouts/readthedocs.org/user_builds/pertpy/envs/529/lib/python3.11/site-packages/ott/__init__.py", line 29, in <module>
    from . import neural
  File "/home/docs/checkouts/readthedocs.org/user_builds/pertpy/envs/529/lib/python3.11/site-packages/ott/neural/__init__.py", line 14, in <module>
    from . import layers, losses, models, solvers
  File "/home/docs/checkouts/readthedocs.org/user_builds/pertpy/envs/529/lib/python3.11/site-packages/ott/neural/models.py", line 19, in <module>
    import optax
  File "/home/docs/checkouts/readthedocs.org/user_builds/pertpy/envs/529/lib/python3.11/site-packages/optax/__init__.py", line 17, in <module>
    from optax import contrib
  File "/home/docs/checkouts/readthedocs.org/user_builds/pertpy/envs/529/lib/python3.11/site-packages/optax/contrib/__init__.py", line 17, in <module>
    from optax.contrib.cocob import cocob
  File "/home/docs/checkouts/readthedocs.org/user_builds/pertpy/envs/529/lib/python3.11/site-packages/optax/contrib/cocob.py", line 26, in <module>
    from optax._src import base
  File "/home/docs/checkouts/readthedocs.org/user_builds/pertpy/envs/529/lib/python3.11/site-packages/optax/_src/base.py", line 19, in <module>
    import chex
  File "/home/docs/checkouts/readthedocs.org/user_builds/pertpy/envs/529/lib/python3.11/site-packages/chex/__init__.py", line 17, in <module>
    from chex._src.asserts import assert_axis_dimension
  File "/home/docs/checkouts/readthedocs.org/user_builds/pertpy/envs/529/lib/python3.11/site-packages/chex/_src/asserts.py", line 26, in <module>
    from chex._src import asserts_internal as _ai
  File "/home/docs/checkouts/readthedocs.org/user_builds/pertpy/envs/529/lib/python3.11/site-packages/chex/_src/asserts_internal.py", line 34, in <module>
    from chex._src import pytypes
  File "/home/docs/checkouts/readthedocs.org/user_builds/pertpy/envs/529/lib/python3.11/site-packages/chex/_src/pytypes.py", line 54, in <module>
    PRNGKey = jax.random.KeyArray
              ^^^^^^^^^^^^^^^^^^^
  File "/home/docs/checkouts/readthedocs.org/user_builds/pertpy/envs/529/lib/python3.11/site-packages/jax/_src/deprecations.py", line 53, in getattr
    raise AttributeError(f"module {module!r} has no attribute {name!r}")
AttributeError: module 'jax.random' has no attribute 'KeyArray'

Version information

No response

[Zethson]

Can be fixed by upgrading all jax dependencies. However, scvi-tools still depends on chex which causes another one.

Need to wait for the latest scvi-tools release

[lihaowhusos]

Some error, holp to solve it in the future

[Zethson]

This issue is because of scvi-tools. Please install scvi-tools from the main branch or wait for its next release (1-2 weeks)

Feb 10, 2024 7:09:14 AM Hao Li @.***>:

Some error, holp to solve it in the future

— Reply to this email directly, view it on GitHub[https://github.com/theislab/pertpy/issues/530#issuecomment-1936894654], or unsubscribe[https://github.com/notifications/unsubscribe-auth/AFHQA2AZKJBLTDGVY4UTUE3YS4FIRAVCNFSM6AAAAABC6NOH2GVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMYTSMZWHA4TINRVGQ]. You are receiving this because you were assigned. [Tracking image][https://github.com/notifications/beacon/AFHQA2FZ7XR4WR6DPSQWPB3YS4FIRA5CNFSM6AAAAABC6NOH2GWGG33NNVSW45C7OR4XAZNMJFZXG5LFINXW23LFNZ2KUY3PNVWWK3TUL5UWJTTTOKVL4.gif]

[Zethson]

Fixed with the latest scvi-tools release