From 462cd27c8c3146c84e3a98dfe638749b79a0f3f2 Mon Sep 17 00:00:00 2001 From: Sait Cakmak Date: Mon, 4 Mar 2024 14:30:33 -0800 Subject: [PATCH] Update docstring for GenerationNode.gen & fit Summary: -- Differential Revision: D54506927 --- ax/modelbridge/generation_node.py | 29 +++++++++++++++++++++-------- 1 file changed, 21 insertions(+), 8 deletions(-) diff --git a/ax/modelbridge/generation_node.py b/ax/modelbridge/generation_node.py index 5697038e6e5..c908d5ba5ad 100644 --- a/ax/modelbridge/generation_node.py +++ b/ax/modelbridge/generation_node.py @@ -256,8 +256,15 @@ def fit( the model kwargs set on each corresponding model spec and the kwargs passed to this method. - NOTE: Local kwargs take precedence over the ones stored in - ``ModelSpec.model_kwargs``. + Args: + experiment: The experiment to fit the model to. + data: The experiment data used to fit the model. + search_space: An optional overwrite for the experiment search space. + optimization_config: An optional overwrite for the experiment + optimization config. + kwargs: Additional keyword arguments to pass to the model's + ``fit`` method. NOTE: Local kwargs take precedence over the ones + stored in ``ModelSpec.model_kwargs``. """ self._model_spec_to_gen_from = None for model_spec in self.model_specs: @@ -283,21 +290,27 @@ def gen( alongside any kwargs passed in to this function (with local kwargs) taking precedent. + NOTE: Models must have been fit prior to calling ``gen``. + NOTE: Some underlying models may ignore the ``n`` argument and produce a + model-determined number of arms. In that case this method will also output + a generator run with number of arms that may differ from ``n``. + Args: - n: Optional nteger representing how many arms should be in the generator + n: Optional integer representing how many arms should be in the generator run produced by this method. When this is ``None``, ``n`` will be determined by the ``ModelSpec`` that we are generating from. pending_observations: A map from metric name to pending observations for that metric, used by some models to avoid resuggesting points that are currently being evaluated. - max_gen_draws_for_deduplication: TODO + max_gen_draws_for_deduplication: Maximum number of attempts for generating + new candidates without duplicates. If non-duplicate candidates are not + generated with these attempts, a ``GenerationStrategyRepeatedPoints`` + exception will be raised. model_gen_kwargs: Keyword arguments, passed through to ``ModelSpec.gen``; these override any pre-specified in ``ModelSpec.model_gen_kwargs``. - NOTE: Models must have been fit prior to calling ``gen``. - NOTE: Some underlying models may ignore the ``n`` argument and produce a - model-determined number of arms. In that case this method will also output - a generator run with number of arms (that can differ from ``n``). + Returns: + A ``GeneratorRun`` containing the newly generated candidates. """ model_spec = self.model_spec_to_gen_from should_generate_run = True