`uv pip install google-vizier[jax]` installs wrong `equinox==0.11.3` version

[152334H]

$ uv venv --seed
Using Python 3.10.12 interpreter at: /usr/bin/python3
Creating virtualenv at: .venv
 + pip==24.2
 + setuptools==75.2.0
 + wheel==0.44.0
Activate with: source .venv/bin/activate
$ . .venv/bin/activate
$ uv pip install google-vizier[jax]
Resolved 47 packages in 256ms
Downloaded 10 packages in 748ms
Installed 46 packages in 18ms
 + absl-py==2.1.0
 + attrs==23.1.0
 + chex==0.1.87
 + cloudpickle==3.1.0
 + decorator==5.1.1
 + dm-tree==0.1.8
 + equinox==0.11.3    # <------ this should be 0.11.7 https://github.com/google/vizier/blob/main/requirements-jax.txt#L9
 + etils==1.10.0
 + flax==0.10.0
 + fsspec==2024.10.0
 + gast==0.6.0
 + google-vizier==0.1.18
 + googleapis-common-protos==1.65.0
...

this is bad because 0.11.3 immediately breaks a basic import:

Traceback (most recent call last):
  File "/server.py", line 2, in <module>
    from vizier.service import servers
  File "/usr/lib/python3.10/site-packages/vizier/service/servers/__init__.py", line 18, in <module>
    from vizier._src.service.vizier_server import DefaultVizierServer
  File "/usr/lib/python3.10/site-packages/vizier/_src/service/vizier_server.py", line 29, in <module>
    from vizier import pythia
  File "/usr/lib/python3.10/site-packages/vizier/pythia.py", line 21, in <module>
    from vizier._src.pythia.local_policy_supporters import InRamPolicySupporter
  File "/usr/lib/python3.10/site-packages/vizier/_src/pythia/local_policy_supporters.py", line 30, in <module>
    from vizier.pyvizier import converters
  File "/usr/lib/python3.10/site-packages/vizier/pyvizier/converters/__init__.py", line 32, in <module>
    from vizier.pyvizier.converters.jnp_converters import PaddedTrialToArrayConverter
  File "/usr/lib/python3.10/site-packages/vizier/pyvizier/converters/jnp_converters.py", line 26, in <module>
    from vizier._src.jax import types as vt
  File "/usr/lib/python3.10/site-packages/vizier/_src/jax/types.py", line 20, in <module>
    import equinox as eqx
  File "/usr/lib/python3.10/site-packages/equinox/__init__.py", line 3, in <module>
    from . import debug as debug, internal as internal, nn as nn
  File "/usr/lib/python3.10/site-packages/equinox/internal/__init__.py", line 46, in <module>
    from ._finalise_jaxpr import (
  File "/usr/lib/python3.10/site-packages/equinox/internal/_finalise_jaxpr.py", line 186, in <module>
    from ._noinline import noinline_p
  File "/usr/lib/python3.10/site-packages/equinox/internal/_noinline.py", line 379, in <module>
    jax.core.pp_eqn_rules[noinline_p] = _noinline_pretty_print
  File "/usr/lib/python3.10/site-packages/jax/_src/deprecations.py", line 52, in getattr
    raise AttributeError(message)
AttributeError: jax.core.pp_eqn_rules was removed in JAX v0.4.34. 

also, equinox==0.11.8 should probably be used instead. see release notes on jax 0.4.33 breaking (vizier installs jax==0.4.35)

[xingyousong]

Hi, thanks for the ping - we're aware of the issue and fixed it in https://github.com/google/vizier/commit/fa018db4715de9036f5db17e293df00955555e6d

But haven't updated the pip package yet. For now, you can install google-vizier from main using our nightly variant: https://pypi.org/project/google-vizier-dev/#history

e.g. pip install google-vizier-dev==0.1.18.dev20241021163849