Importing jax-moseq produces deprecation error

[wingillis]

A fresh install of jax-moseq installs jax==0.4.20, which deprecates some classes that chex==0.1.6 depends on, and produces the following error upon import:
AttributeError: module 'jax.interpreters.pxla' has no attribute 'ShardedDeviceArray'

The solution is to update the version of chex to the latest (0.1.84), which no longer references ShardedDeviceArray.

[calebweinreb]

So are you saying that chex should be pinned to 0.1.84?

[wingillis]

Yes

[calebweinreb]

I'm a little worried about this given that keypoint-moseq install currently works for most people, and upgrading chex/jax might open a whole can of worms. When you say "a fresh install of jax-moseq installs jax==0.4.20", what do you mean exactly? Like if you update jax-moseq it tries to update jax too?

[bainro]

google-deepmind/chex#251
https://github.com/google-deepmind/chex/tags

Merged to chex master in march, so probably 0.1.8.

I'm having the exact same issue on Ubuntu 22 with GPU. Installing a newer chex gives dependency error from jax-moseq.

[calebweinreb]

Is the dependency error in pip? If so, I think it can just be ignored.

[bainro]

Yeah, just edited my comment as you replied. I did ignore it, thanks Caleb :)