8000 algorithms/designers/scalarization.py: modernize jaxtyping by ConnorBaker · Pull Request #1029 · google/vizier · GitHub
[go: up one dir, main page]
Skip to content

algorithms/designers/scalarization.py: modernize jaxtyping - #1029

Closed
ConnorBaker wants to merge 1 commit into
google:mainfrom
ConnorBaker:fix/modernize-jaxtyping
Closed

algorithms/designers/scalarization.py: modernize jaxtyping#1029
ConnorBaker wants to merge 1 commit into
google:mainfrom
ConnorBaker:fix/modernize-jaxtyping

Conversation

@ConnorBaker
@ConnorBaker ConnorBaker commented Jan 9, 2024
Copy link
Copy Markdown

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
ConnorBaker commented Jan 9, 2024
8000 Copy link
Copy Markdown
Author

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

@ConnorBaker ConnorBaker mentioned this pull request Jan 9, 2024
Closed
copybara-service Bot pushed a commit that referenced this pull request Jan 9, 2024
@sagipe @copybara-github
Apply suggestion from ConnorBaker in #1029
a08e5f6 
PiperOrigin-RevId: 596802885
@copybara-service copybara-service Bot mentioned this pull request Jan 9, 2024
Merged
copybara-service Bot pushed a commit that referenced this pull request Jan 9, 2024
@sagipe @copybara-github
Apply suggestion from ConnorBaker in #1029
eb27ec9 
PiperOrigin-RevId: 596802885
copybara-service Bot pushed a commit that referenced this pull request Jan 9, 2024
@sagipe @copybara-github
Apply suggestion from ConnorBaker in #1029
2877348 
PiperOrigin-RevId: 596925777
@sagipe
sagipe commented Jan 9, 2024
Copy link
Copy Markdown
Member

Thanks Connor, applied in #1033
Best,
Sagi

@ConnorBaker ConnorBaker closed this Jan 9, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@ConnorBaker @sagipe
0