Incompatibility between VizierSampler and JAX ≥ 0.4.35

[andresckamilo]

Expected behavior

I was using the VizierSampler from OptunaHub, and when starting the optimization, the process failed due to a compatibility issue between JAX and Equinox (AttributeError: module ‘jax.core’ has no attribute ‘Primitive’).

Environment

• The package of OptunaHub which you are using:

• Optuna version: 4.5.0

• OptunaHub version: 0.3.1

• Python version: 3.13

• OS: MacOS

Other libraries:

• jax==0.4.35
• jaxlib==0.4.35
• equinox==0.11.2
• flax==0.7.5
• xgboost==2.1.1
• scikit-learn==1.5.2

Error messages, stack traces, or logs

55   warnings.warn(message, DeprecationWarning, stacklevel=2)
     56   return fn
---> 57 raise AttributeError(f"module {module!r} has no attribute {name!r}")

AttributeError: module 'jax.core' has no attribute 'Primitive'

Steps to reproduce

import optuna
import optunahub
from xgboost import XGBClassifier
from sklearn.model_selection import cross_val_score, StratifiedKFold
import numpy as np

X = np.random.rand(100, 10)
y = np.random.randint(0, 2, 100)

vizier = optunahub.load_module("samplers/vizier")

def objective(trial):
max_depth = trial.suggest_int("max_depth", 2, 10)
learning_rate = trial.suggest_float("learning_rate", 0.01, 0.3)
subsample = trial.suggest_float("subsample", 0.6, 0.9)
n_estimators = trial.suggest_int("n_estimators", 100, 300)

model = XGBClassifier(
    objective="binary:logistic",
    eval_metric="auc",
    max_depth=max_depth,
    learning_rate=learning_rate,
    subsample=subsample,
    n_estimators=n_estimators,
    use_label_encoder=False,
)

cv = StratifiedKFold(n_splits=3, shuffle=True, random_state=42)
scores = cross_val_score(model, X, y, cv=cv, scoring="roc_auc")
return scores.mean()

study = optuna.create_study(
study_name="vizier_test",
direction="maximize",
sampler=vizier.VizierSampler()
)

study.optimize(objective, n_trials=10)

Additional context (optional)

No response