brentyi / jaxlie

Rigid transforms + Lie groups in JAX
MIT License
209 stars 15 forks source link

jax.random.KeyArray is deprecated #13

Closed mjo22 closed 6 months ago

mjo22 commented 7 months ago

Hello! I have been using jaxlie for a project of mine, see here. I wanted to report a depreciation warning

/path/to/jaxlie/hints/__init__.py:25
  /path/to/jaxlie/hints/__init__.py:25: DeprecationWarning: jax.random.KeyArray is deprecated. Use jax.Array for annotations, and jax.dtypes.issubdtype(arr.dtype, jax.dtypes.prng_key) for runtime detection of typed prng keys (i.e. keys created with jax.random.key).
  For more information, see https://jax.readthedocs.io/en/latest/jep/9263-typed-keys.html
    from jax.random import KeyArray

See https://jax.readthedocs.io/en/latest/changelog.html#jax-0-4-16-sept-18-2023 for the changelog. If this could be updated on the next release that would be great. Thank you for writing a great package!

brentyi commented 6 months ago

Just did a release that should fix this. Thanks for reporting!