Farama-Foundation / Farama-Foundation/Jumpy
Update to use `jax.Array`
Open
- Dominant language
- Python
- Stars
- 56
- Forks
- 9
- PR merge metrics
- No merged PRs in 30d
Description
In `_which_np`, we check if the variable is a `jnp.DeviceArray` however this will not work for `ShardedDeviceArray` and `GlobalDeviceArray`.
In Jax [0.4](https://github.com/google/jax/releases/tag/jax-v0.4.1), this is changed to `jax.Array` as a unified array type, we could either add support for the other array types not currently supported. Or we update to use jax 0.4
Contributor guide
No contributing guide indexed for this repository
Assessment
This issue has not been assessed yet.