Skip to content

Commit

Permalink
use generation_node_name on the GeneratorRun (facebook#2003)
Browse files Browse the repository at this point in the history
Summary:

Replace the use of generation_step_index with generation_node_name because GenerationStrategy should no longer access index

Differential Revision: D51432075
  • Loading branch information
Mia Garrard authored and facebook-github-bot committed Nov 29, 2023
1 parent e6ea088 commit acf6940
Showing 1 changed file with 9 additions and 5 deletions.
14 changes: 9 additions & 5 deletions ax/modelbridge/generation_strategy.py
Original file line number Diff line number Diff line change
Expand Up @@ -508,11 +508,15 @@ def _get_model_state_from_last_generator_run(self) -> Dict[str, Any]:
# split them per-model there.
model_state_on_lgr = {}
model_on_curr = self._curr.model_enum
if (
lgr is not None
and lgr._generation_step_index == self._curr.index
and lgr._model_state_after_gen
):
if lgr is None:
return model_state_on_lgr

if all(isinstance(s, GenerationStep) for s in self._steps):
grs_equal = lgr._generation_step_index == self._curr.index
else:
grs_equal = lgr._generation_node_name == self._curr.node_name

if grs_equal and lgr._model_state_after_gen:
if self.model or isinstance(model_on_curr, ModelRegistryBase):
# TODO[drfreund]: Consider moving this to `GenerationStep` or
# `GenerationNode`.
Expand Down

0 comments on commit acf6940

Please sign in to comment.