algorithms/designers/scalarization.py: modernize jaxtyping

[ConnorBaker]

Carrying on the work done in f7e79e1.

Without this change, the following (using the latest release of Vizier on PyPi with the JAX extra) code produces the warnings below.

Code

from vizier import service
from vizier.service import clients
from vizier.service import pyvizier as vz


def evaluate(x: float, y: float) -> float:
    return x**2 - y**2


if __name__ == "__main__":
    problem = vz.ProblemStatement()
    _ = problem.search_space.root.add_float_param("x", 0.0, 1.0)
    _ = problem.search_space.root.add_float_param("y", 0.0, 1.0)
    problem.metric_information.append(
        vz.MetricInformation(name="maximize_metric", goal=vz.ObjectiveMetricGoal.MAXIMIZE)
    )

    study_config = vz.StudyConfig.from_problem(problem)
    study_config.algorithm = "GAUSSIAN_PROCESS_BANDIT"

    study_client = clients.Study.from_study_config(study_config, owner="owner", study_id="example_study_id")
    print("Local SQL database file located at: ", service.VIZIER_DB_PATH)

    for i in range(10):
        for suggestion in study_client.suggest(count=1):
            x = suggestion.parameters["x"]
            y = suggestion.parameters["y"]
            objective = evaluate(x, y)
            print(f"Iteration {i}, suggestion ({x},{y}) led to objective value {objective}.")
            final_measurement = vz.Measurement({"maximize_metric": objective})
            _ = suggestion.complete(final_measurement)

    for optimal_trial in map(lambda trial: trial.materialize(), study_client.optimal_trials()):
        print(
            "Optimal Trial Suggestion and Objective:",
            optimal_trial.parameters,
            optimal_trial.final_measurement,
        )

Output

WARNING:absl:Python 3.8+ is required in this case.
Local SQL database file located at:  /Users/connorbaker/micromamba/envs/ghc_hyperopt/lib/python3.11/site-packages/vizier/_src/service/vizier.db
/Users/connorbaker/micromamba/envs/ghc_hyperopt/lib/python3.11/site-packages/vizier/_src/algorithms/designers/scalarization.py:38: UserWarning: As of jaxtyping version 0.2.24, jaxtyping now prefers the syntax
```
from jaxtyping import jaxtyped
# Use your favourite typechecker: usually one of the two lines below.
from typeguard import typechecked as typechecker
from beartype import beartype as typechecker

@jaxtyped(typechecker=typechecker)
def foo(...):
```
and the old double-decorator syntax
```
@jaxtyped
@typechecker
def foo(...):
```
should no longer be used. (It will continue to work as it did before, but the new approach will produce more readable error messages.)
In particular note that `typechecker` must be passed via keyword argument; the following is not valid:
```
@jaxtyped(typechecker)
def foo(...):
```

  @jt.jaxtyped
/Users/connorbaker/micromamba/envs/ghc_hyperopt/lib/python3.11/site-packages/vizier/_src/algorithms/designers/scalarization.py:55: UserWarning: As of jaxtyping version 0.2.24, jaxtyping now prefers the syntax
```
from jaxtyping import jaxtyped
# Use your favourite typechecker: usually one of the two lines below.
from typeguard import typechecked as typechecker
from beartype import beartype as typechecker

@jaxtyped(typechecker=typechecker)
def foo(...):
```
and the old double-decorator syntax
```
@jaxtyped
@typechecker
def foo(...):
```
should no longer be used. (It will continue to work as it did before, but the new approach will produce more readable error messages.)
In particular note that `typechecker` must be passed via keyword argument; the following is not valid:
```
@jaxtyped(typechecker)
def foo(...):
```

  @jt.jaxtyped
/Users/connorbaker/micromamba/envs/ghc_hyperopt/lib/python3.11/site-packages/vizier/_src/algorithms/designers/scalarization.py:66: UserWarning: As of jaxtyping version 0.2.24, jaxtyping now prefers the syntax
```
from jaxtyping import jaxtyped
# Use your favourite typechecker: usually one of the two lines below.
from typeguard import typechecked as typechecker
from beartype import beartype as typechecker

@jaxtyped(typechecker=typechecker)
def foo(...):
```
and the old double-decorator syntax
```
@jaxtyped
@typechecker
def foo(...):
```
should no longer be used. (It will continue to work as it did before, but the new approach will produce more readable error messages.)
In particular note that `typechecker` must be passed via keyword argument; the following is not valid:
```
@jaxtyped(typechecker)
def foo(...):
```

  @jt.jaxtyped
/Users/connorbaker/micromamba/envs/ghc_hyperopt/lib/python3.11/site-packages/vizier/_src/algorithms/designers/scalarization.py:81: UserWarning: As of jaxtyping version 0.2.24, jaxtyping now prefers the syntax
```
from jaxtyping import jaxtyped
# Use your favourite typechecker: usually one of the two lines below.
from typeguard import typechecked as typechecker
from beartype import beartype as typechecker

@jaxtyped(typechecker=typechecker)
def foo(...):
```
and the old double-decorator syntax
```
@jaxtyped
@typechecker
def foo(...):
```
should no longer be used. (It will continue to work as it did before, but the new approach will produce more readable error messages.)
In particular note that `typechecker` must be passed via keyword argument; the following is not valid:
```
@jaxtyped(typechecker)
def foo(...):
```

  @jt.jaxtyped
/Users/connorbaker/micromamba/envs/ghc_hyperopt/lib/python3.11/site-packages/vizier/_src/algorithms/designers/scalarization.py:109: UserWarning: As of jaxtyping version 0.2.24, jaxtyping now prefers the syntax
```
from jaxtyping import jaxtyped
# Use your favourite typechecker: usually one of the two lines below.
from typeguard import typechecked as typechecker
from beartype import beartype as typechecker

@jaxtyped(typechecker=typechecker)
def foo(...):
```
and the old double-decorator syntax
```
@jaxtyped
@typechecker
def foo(...):
```
should no longer be used. (It will continue to work as it did before, but the new approach will produce more readable error messages.)
In particular note that `typechecker` must be passed via keyword argument; the following is not valid:
```
@jaxtyped(typechecker)
def foo(...):
```

  @jt.jaxtyped
Iteration 0, suggestion (0.5,0.5) led to objective value 0.0.
Iteration 1, suggestion (0.7199959629810866,0.28000490122422056) led to objective value 0.4399914419994767.
Iteration 2, suggestion (0.9599940200424931,0.04000513419520845) led to objective value 0.9199881077553701.
Iteration 3, suggestion (0.7838943928094828,0.0) led to objective value 0.6144904190781477.
Iteration 4, suggestion (0.9999954768912652,0.5599887043821371) led to objective value 0.6864036047674045.
Iteration 5, suggestion (1.0,0.0) led to objective value 1.0.
Iteration 6, suggestion (0.9999999999999976,2.6172520215590815e-15) led to objective value 0.9999999999999951.
Iteration 7, suggestion (1.0,0.03269750221434778) led to objective value 0.9989308733489427.
Iteration 8, suggestion (0.1400016244462224,0.14004559082072232) led to objective value -1.231266074408141e-05.
Iteration 9, suggestion (1.0,0.010637668977110257) led to objective value 0.9998868399987334.
Optimal Trial Suggestion and Objective: ParameterDict(_items={'x': 1.0, 'y': 0.0}) Measurement(metrics={'maximize_metric': Metric(value=1.0, std=None)}, elapsed_secs=0.0, steps=0, checkpoint_path='')

[ConnorBaker]

CC @sagipe as this is basically a copy-paste of the change you made previously.

[sagipe]

Thanks Connor, applied in #1033
Best,
Sagi