search

inseq-team / inseq

Interpretability for sequence generation models 🐛 🔍
https://inseq.org
Apache License 2.0
375 stars 36 forks source link

Module 'jax' has no attribute 'typing' #234

Closed carschno closed 12 months ago

carschno commented 1 year ago

🐛 Bug Report

As reported in https://github.com/inseq-team/inseq/issues/232#issuecomment-1807001439, there might be an issue related to jaxtyping (#167).

AttributeError                            Traceback (most recent call last)
Cell In [1], line 1
----> 1 import inseq

[...]

File /usr/local/lib/python3.9/dist-packages/jaxtyping/_array_types.py:680
    678 PRNGKeyArray = Union[Key[jax.Array, ""], UInt32[jax.Array, "2"]]
    679 Scalar = Shaped[jax.Array, ""]
--> 680 ScalarLike = Shaped[jax.typing.ArrayLike, ""]

AttributeError: module 'jax' has no attribute 'typing'

It looks like updating jax and/or jaxlib has resolved the issue: https://github.com/inseq-team/inseq/issues/232#issuecomment-1807043699

Similar issues in other projects:

🔬 How To Reproduce

Steps to reproduce the behavior:

Unclear.

Code sample

import inseq

Environment

python --version

Expected behavior

Import Inseq without errors, with or without jax/jaxlib being installed.

Additional context

Currently, Jax is not a dependency of Inseq. However, many users will have Jax installed. We need to make sure that Inseq works and possibly document/pin Jax versions that are compatible.

carschno commented 12 months ago

I have not yet been able to reproduce the issue. @seemavishal would you be able to provide the Jax/Jaxlib versions you had installed before updating them (which apparently resolved the issue)?

carschno commented 12 months ago

Trying different combinations of jax and jaxlib versions, all with Python 3.9. Checking whether import inseq works:

status
jax not installed 0.3.0 0.3.0 0.4.0 0.4.20 0.4.20 0.4.20 not installed not installed
jaxlib not installed not installed 0.4.0 0.4.0 0.4.0 0.4.20 not installed 0.4.20 0.4.0
works ✔️ ✔️ ❌ [^1] ❌ [^2] ❌ [^3] ✔️ ✔️ ✔️ ✔️

[^1]: RuntimeError: jaxlib version 0.4.0 is newer than and incompatible with jax version 0.3.0 [^2]: AttributeError: module 'jax' has no attribute 'typing' [^3]: RuntimeError: jaxlib is version 0.4.0, but this version of jax requires version >= 0.4.14.

We are interested in error case 2 here, so it looks like a recent version of both jax and jaxlib is required.

SVC04 commented 12 months ago

Hi,

To resolve the error related to jaxlib I tried pip install -U jaxlib.

That resolved the problem.

On Mon, Nov 13, 2023, 4:45 PM Carsten Schnober @.***> wrote:

Trying different combinations of jax and jaxlib versions, all with Python 3.9. Checking whether import inseq works: status jax not installed 0.3.0 0.3.0 0.4.0 0.4.20 0.4.20 jaxlib not installed not installed 0.4.0 0.4.0 0.4.0 0.4.20 works ✔️ ✔️ ❌ 1 <#m_-2254177894444722280m-6800094939708900059user-content-fn-1-a82e3d774e3fdc57498976e403fa3a9a> ❌ 2 <#m-2254177894444722280m-6800094939708900059user-content-fn-2-a82e3d774e3fdc57498976e403fa3a9a> ❌ 3 <#m-2254177894444722280m-6800094939708900059_user-content-fn-3-a82e3d774e3fdc57498976e403fa3a9a> ✔️

We are interested in error case 2 here, so it looks like a recent version of both jax and jaxlib is required. Footnotes

1.

RuntimeError: jaxlib version 0.4.0 is newer than and incompatible with jax version 0.3.0 ↩ <#m_-2254177894444722280m-6800094939708900059_user-content-fnref-1-a82e3d774e3fdc57498976e403fa3a9a> 2.

AttributeError: module 'jax' has no attribute 'typing' ↩ <#m_-2254177894444722280m-6800094939708900059_user-content-fnref-2-a82e3d774e3fdc57498976e403fa3a9a> 3.

RuntimeError: jaxlib is version 0.4.0, but this version of jax requires version >= 0.4.14. ↩ <#m_-2254177894444722280m-6800094939708900059_user-content-fnref-3-a82e3d774e3fdc57498976e403fa3a9a>

— Reply to this email directly, view it on GitHub https://github.com/inseq-team/inseq/issues/234#issuecomment-1808099069, or unsubscribe https://github.com/notifications/unsubscribe-auth/AU7Z435WZZYYUY7DJVSSEYLYEII7FAVCNFSM6AAAAAA7H4PL4SVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMYTQMBYGA4TSMBWHE . You are receiving this because you were mentioned.Message ID: @.***>