Skip to content

No-op cleanups. #1281

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
copybara-service merged 1 commit into from
Sep 12, 2025
Merged

No-op cleanups. #1281

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
84 changes: 46 additions & 38 deletions vizier/pyvizier/converters/spatio_temporal.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@
import collections
import copy
import dataclasses
from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
from typing import Any, Literal, Optional, Sequence, Union

from absl import logging
import numpy as np
Expand All @@ -37,7 +37,7 @@ class TimedLabels:
labels: (M, 1) arrays keyed by strings, corresponding to metrics.
"""
times: np.ndarray
labels: Dict[str, np.ndarray]
labels: dict[str, np.ndarray]


class TimedLabelsExtractor:
Expand All @@ -48,12 +48,14 @@ class TimedLabelsExtractor:
CUMMAX_FIRSTONLY = 'cummax_firstonly'
RAW = 'raw'

def __init__(self,
metric_converters: Sequence[core.ModelOutputConverter],
timestamp: str = 'steps',
*,
temporal_index_points: Sequence[float] = tuple(),
value_extraction: str = 'cummax_lastonly'):
def __init__(
self,
metric_converters: Sequence[core.ModelOutputConverter],
timestamp: Literal['steps', 'elapsed_secs', 'index'] = 'steps',
*,
temporal_index_points: Sequence[float] = tuple(),
value_extraction: str = 'cummax_lastonly',
):
"""Init.

Args:
Expand Down Expand Up @@ -88,6 +90,8 @@ def __init__(self,
self.CUMMAX_FIRSTONLY):
raise ValueError(
'Bad value for value_extraction rule: {}'.format(value_extraction))
if timestamp not in ('steps', 'elapsed_secs', 'index'):
raise ValueError(f'Invalid timestamp: {timestamp}')
if value_extraction in [self.CUMMAX_LASTONLY, self.CUMMAX_FIRSTONLY]:
if len(metric_converters) > 1:
raise ValueError(
Expand All @@ -104,7 +108,7 @@ def _cummax_fn(self, metric_converter: core.ModelOutputConverter) -> Any:
else:
return np.minimum

def convert(self, trials: Sequence[pyvizier.Trial]) -> List[TimedLabels]:
def convert(self, trials: Sequence[pyvizier.Trial]) -> list[TimedLabels]:
"""Converts each trial into TimedLabels object."""
timedlabels = []
if self.temporal_index_points.size == 0:
Expand Down Expand Up @@ -193,15 +197,15 @@ def convert(self, trials: Sequence[pyvizier.Trial]) -> List[TimedLabels]:
def to_timestamps(self,
measurements: Sequence[pyvizier.Measurement]) -> np.ndarray:
""""Returns an arry of shape (len(measurements), 1)."""
timestamps = []

for idx, measurement in enumerate(measurements):
if self.timestamp == 'steps':
timestamps.append(measurement.steps)
if self.timestamp == 'elapsed_secs':
timestamps.append(measurement.elapsed_secs)
if self.timestamp == 'index':
timestamps.append(idx)
if self.timestamp == 'steps':
timestamps = [m.steps for m in measurements]
elif self.timestamp == 'elapsed_secs':
timestamps = [m.elapsed_secs for m in measurements]
elif self.timestamp == 'index':
timestamps = list(range(len(measurements)))
else:
# Should be unreachable given check in __init__.
raise ValueError(f'Invalid timestamp: {self.timestamp}')
return np.asarray(timestamps)[:, np.newaxis]

def extract_all_timestamps(
Expand Down Expand Up @@ -245,7 +249,7 @@ def __init__(self,

def to_xy(
self, trials: Sequence[pyvizier.Trial]
) -> Tuple[Dict[str, np.ndarray], Dict[str, np.ndarray]]:
) -> tuple[dict[str, np.ndarray], dict[str, np.ndarray]]:
"""Returned values can be used as `x` and `y` for keras.Model.fit().

Args:
Expand Down Expand Up @@ -280,7 +284,7 @@ def to_xy(

return all_features, all_labels

def to_features(self, trial, temporal_index_points) -> Dict[str, np.ndarray]:
def to_features(self, trial, temporal_index_points) -> dict[str, np.ndarray]:
"""Converts a trial at the specified index points.

Args:
Expand All @@ -306,14 +310,14 @@ def to_features(self, trial, temporal_index_points) -> Dict[str, np.ndarray]:
return features

@property
def features_shape(self) -> Dict[str, Sequence[Union[int, None]]]:
def features_shape(self) -> dict[str, Sequence[Union[int, None]]]:
"""Returned value can be used as `input_shape` for keras.Model.build()."""
shapes = copy.deepcopy(self.trial_converter.features_shape)
shapes[self.timed_labels_extractor.timestamp] = (None, 1)
return shapes

@property
def output_specs(self) -> Dict[str, core.NumpyArraySpec]:
def output_specs(self) -> dict[str, core.NumpyArraySpec]:
specs = copy.deepcopy(self.trial_converter.output_specs)
name = self.timed_labels_extractor.timestamp
# Can't use float32 max, because
Expand All @@ -325,11 +329,11 @@ def output_specs(self) -> Dict[str, core.NumpyArraySpec]:
return specs

@property
def labels_shape(self) -> Dict[str, Sequence[Union[int, None]]]:
def labels_shape(self) -> dict[str, Sequence[Union[int, None]]]:
return self.trial_converter.labels_shape

@property
def metric_information(self) -> Dict[str, pyvizier.MetricInformation]:
def metric_information(self) -> dict[str, pyvizier.MetricInformation]:
return self.trial_converter.metric_information


Expand Down Expand Up @@ -358,8 +362,8 @@ def __init__(self,
self.temporal_index_points = temporal_index_points

def _single_timedlabels_to_temporal_observations(
self, timed_labels: TimedLabels,
ts: Union[Sequence[float], Sequence[int]]) -> Dict[str, np.ndarray]:
self, timed_labels: TimedLabels, ts: Union[Sequence[float], Sequence[int]]
) -> dict[str, np.ndarray]:
"""Subroutine of _to_temporal_observations().

Args:
Expand Down Expand Up @@ -409,8 +413,10 @@ def _single_timedlabels_to_temporal_observations(
return {k: np.asarray(v) for k, v in this_labels.items()}

def _to_temporal_observations(
self, timed_labels_sequence: Sequence[TimedLabels],
ts: Union[Sequence[float], Sequence[int]]) -> Dict[str, np.ndarray]:
self,
timed_labels_sequence: Sequence[TimedLabels],
ts: Union[Sequence[float], Sequence[int]],
) -> dict[str, np.ndarray]:
"""Returns a dict of np arrays of temporal observations.

Filters timed_labels.labels to leave only the time indices that appear in
Expand All @@ -437,7 +443,7 @@ def _to_temporal_observations(

def to_xy(
self, trials: Sequence[pyvizier.Trial]
) -> Tuple[Dict[str, np.ndarray], Dict[str, np.ndarray]]:
) -> tuple[dict[str, np.ndarray], dict[str, np.ndarray]]:
"""Returned value can be used as `x`, `y` for keras.Model.fit()."""
all_timed_labels = self.timed_labels_extractor.convert(trials)
labels = self._to_temporal_observations(all_timed_labels,
Expand All @@ -446,33 +452,34 @@ def to_xy(
return self.trial_converter.to_features(trials), labels

@property
def features_shape(self) -> Dict[str, Sequence[Union[int, None]]]:
def features_shape(self) -> dict[str, Sequence[Union[int, None]]]:
"""Returned value can be used as `input_shape` for keras.Model.build()."""
return self.trial_converter.features_shape

@property
def output_specs(self) -> Dict[str, core.NumpyArraySpec]:
def output_specs(self) -> dict[str, core.NumpyArraySpec]:
return self.trial_converter.output_specs

@property
def labels_shape(self) -> Dict[str, Sequence[Union[int, None]]]:
def labels_shape(self) -> dict[str, Sequence[Union[int, None]]]:
shapes = dict()
for mc in self.trial_converter.metric_converters:
shapes[mc.metric_information.name] = (None,
len(self.temporal_index_points))
return shapes

@property
def metric_information(self) -> Dict[str, pyvizier.MetricInformation]:
def metric_information(self) -> dict[str, pyvizier.MetricInformation]:
return self.trial_converter.metric_information

def to_features(self,
trials: Sequence[pyvizier.Trial]) -> Dict[str, np.ndarray]:
def to_features(
self, trials: Sequence[pyvizier.Trial]
) -> dict[str, np.ndarray]:
return self.trial_converter.to_features(trials)

def to_xty(
self, trials: Sequence[pyvizier.Trial], temporal_selection: str = 'auto'
) -> Tuple[Dict[str, np.ndarray], np.ndarray, Dict[str, np.ndarray]]:
) -> tuple[dict[str, np.ndarray], np.ndarray, dict[str, np.ndarray]]:
"""Returns x, t, and y.

Args:
Expand All @@ -490,8 +497,9 @@ def to_xty(
observations: Dict of length equal to metrics, whose values
have shape [len(trials), len(temporal_index_points)]. May contain NaNs.
"""
timed_labels: List[TimedLabels] = (
self.timed_labels_extractor.convert(trials))
timed_labels: list[TimedLabels] = self.timed_labels_extractor.convert(
trials
)

if temporal_selection == 'default' or (temporal_selection == 'auto' and
self.temporal_index_points.size):
Expand Down