Fix the haiku example when running with jax 0.5.3 and jaxlib 0.6.0 #319
Description
Expected behavior
After the release of jax 0.5.3 and jaxlib 0.6.0, the CI related to haiku started failing. To temporarily resolve the issue, I opened a PR to add the version constraint. This issue proposes resolving the failure.
Environment
See https://github.com/optuna/optuna-examples/actions/runs/14575803940
Error messages, stack traces, or logs
Traceback (most recent call last):
File "/home/runner/work/optuna-examples/optuna-examples/haiku/haiku_simple.py", line 25, in <module>
import haiku as hk
File "/opt/hostedtoolcache/Python/3.11.12/x64/lib/python3.11/site-packages/haiku/__init__.py", line 20, in <module>
from haiku import experimental
File "/opt/hostedtoolcache/Python/3.11.12/x64/lib/python3.11/site-packages/haiku/experimental/__init__.py", line 55, in <module>
from haiku.experimental import jaxpr_info
File "/opt/hostedtoolcache/Python/3.11.12/x64/lib/python3.11/site-packages/haiku/experimental/jaxpr_info.py", line [18](https://github.com/optuna/optuna-examples/actions/runs/14575803940/job/40881567302#step:5:19), in <module>
from haiku._src.jaxpr_info import as_html
File "/opt/hostedtoolcache/Python/3.11.12/x64/lib/python3.11/site-packages/haiku/_src/jaxpr_info.py", line 101, in <module>
ComputeFlopsFn = Callable[[jax.core.JaxprEqn, Expression], int]
^^^^^^^^^^^^^^^^^
File "/opt/hostedtoolcache/Python/3.11.12/x64/lib/python3.11/site-packages/jax/_src/deprecations.py", line 54, in getattr
raise AttributeError(message)
AttributeError: jax.core.JaxprEqn was removed in JAX v0.6.0. Use jax.extend.core.JaxprEqn instead, and see https://docs.jax.dev/en/latest/jax.extend.html for details.
Steps to reproduce
Run haiku example with jax 0.5.3 and jaxlib 0.6.0
Reproducible examples (optional)
# python code