BUG: Fix is_jax_array for jax>=0.8.2
#369
Merged
Loading
is_jax_array for jax>=0.8.2
#369