Skip to content

Commit

Permalink
Update GenerationNode.model_to_gen_from_name (#2407)
Browse files Browse the repository at this point in the history
Summary:
Pull Request resolved: #2407

Updates the behavior to match the docstring. Previously, this would call `model_spec_to_gen_from`, which could lead to fitting multiple models if the node had multiple model specs. The return type was always `str` rather than `Optional[str]` suggested by the type hints.

Reviewed By: mgarrard

Differential Revision: D56735619

fbshipit-source-id: a6b39b3bc921f720bcb43daece55dc9da1ee96c6
  • Loading branch information
saitcakmak authored and facebook-github-bot committed May 2, 2024
1 parent c5f50e7 commit 9ba941c
Show file tree
Hide file tree
Showing 2 changed files with 6 additions and 1 deletion.
5 changes: 4 additions & 1 deletion ax/modelbridge/generation_node.py
Original file line number Diff line number Diff line change
Expand Up @@ -143,7 +143,10 @@ def model_to_gen_from_name(self) -> Optional[str]:
"""Returns the name of the model that will be used for gen, if there is one.
Otherwise, returns None.
"""
return self.model_spec_to_gen_from.model_key
if self._model_spec_to_gen_from is not None:
return self._model_spec_to_gen_from.model_key
else:
return None

@property
def generation_strategy(self) -> modelbridge.generation_strategy.GenerationStrategy:
Expand Down
2 changes: 2 additions & 0 deletions ax/modelbridge/tests/test_generation_node.py
Original file line number Diff line number Diff line change
Expand Up @@ -113,6 +113,7 @@ def test_properties(self) -> None:
),
],
)
self.assertIsNone(node.model_to_gen_from_name)
dat = self.branin_experiment.lookup_data()
node.fit(
experiment=self.branin_experiment,
Expand All @@ -124,6 +125,7 @@ def test_properties(self) -> None:
self.assertEqual(
node.model_spec_to_gen_from.model_kwargs, node.model_specs[0].model_kwargs
)
self.assertEqual(node.model_to_gen_from_name, "GPEI")
self.assertEqual(
node.model_spec_to_gen_from.model_gen_kwargs,
node.model_specs[0].model_gen_kwargs,
Expand Down

0 comments on commit 9ba941c

Please sign in to comment.