Broken nightly testing due to breaking changes in JAX v0.4.36

**CI job:** https://github.com/daskol/yax/actions/runs/12194083414/job/34017355642.
**Upstream discussion:** https://github.com/jax-ml/jax/issues/25314.

Since JAX v0.4.36, `jax.core.ConcreteArray` is not a part of public API.

```python
yax.py:40: in <module>
    from jax.core import (
ImportError: cannot import name 'ConcreteArray' from 'jax.core'
```