Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Cannot import oryx due to jax AttributeError #108

@genos

Description

@genos

First off, thanks very much for oryx, it seems really interesting! Unfortunately, when I tried to play around with it locally, I ran into an AttributeError due, I think, to falling behind jax versions:

>>> import oryx
Traceback (most recent call last):
  File "<python-input-0>", line 1, in <module>
    import oryx
  File "/Users/graham/tmp/reprex/.venv/lib/python3.13/site-packages/oryx/__init__.py", line 17, in <module>
    from oryx import core
  File "/Users/graham/tmp/reprex/.venv/lib/python3.13/site-packages/oryx/core/__init__.py", line 16, in <module>
    from oryx.core import interpreters
  File "/Users/graham/tmp/reprex/.venv/lib/python3.13/site-packages/oryx/core/interpreters/__init__.py", line 16, in <module>
    from oryx.core.interpreters import harvest
  File "/Users/graham/tmp/reprex/.venv/lib/python3.13/site-packages/oryx/core/interpreters/harvest.py", line 1568, in <module>
    @shard_map.register_check(sow_p)
     ^^^^^^^^^^^^^^^^^^^^^^^^
AttributeError: module 'jax.experimental.shard_map' has no attribute 'register_check'

Regarding versions, here's the output of setting up a new project with uv and calling uv add oryx:

Click to expand
Using CPython 3.13.3
Creating virtual environment at: .venv
Resolved 17 packages in 4ms
Installed 16 packages in 79ms
 + absl-py==2.3.0
 + attrs==25.3.0
 + cloudpickle==3.1.1
 + decorator==5.2.1
 + dm-tree==0.1.9
 + gast==0.6.0
 + jax==0.6.2
 + jaxlib==0.6.2
 + ml-dtypes==0.5.1
 + numpy==2.3.1
 + opt-einsum==3.4.0
 + oryx==0.2.9
 + scipy==1.16.0
 + six==1.17.0
 + tensorflow-probability==0.25.0
 + wrapt==1.17.2

Thanks again!

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type
    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions