docstr: Tests that it is possible to restore a model from a generator run it
  produced, if `Models` registry was used.
code: def test_get_model_from_generator_run(self) -> None:
      """Tests that it is possible to restore a model from a generator run it
      produced, if `Models` registry was used.
      """
      exp = get_branin_experiment()
      initial_sobol = Models.SOBOL(experiment=exp, seed=239)
      gr = initial_sobol.gen(n=1)
      # Restore the model as it was before generation.
      sobol = get_model_from_generator_run(
          generator_run=gr,
          experiment=exp,
          data=exp.fetch_data(),
          models_enum=Models,
          after_gen=False,
      )
      self.assertEqual(sobol.model.init_position, 0)
      self.assertEqual(sobol.model.seed, 239)
      # Restore the model as it was after generation (to resume generation).
      sobol_after_gen = get_model_from_generator_run(
          generator_run=gr,
          experiment=exp,
          data=exp.fetch_data(),
          models_enum=Models,
      )
      self.assertEqual(sobol_after_gen.model.init_position, 1)
      self.assertEqual(sobol_after_gen.model.seed, 239)
      self.assertEqual(
          initial_sobol.gen(n=1).arms, sobol_after_gen.gen(n=1).arms
      )
      exp.new_trial(generator_run=gr)
      # Check restoration of GPEI, to ensure proper restoration of callable kwargs
      gpei = Models.GPEI(experiment=exp, data=get_branin_data())
      # Punch GPEI model + bridge kwargs into the Sobol generator run, to avoid
      # a slow call to `gpei.gen`, and remove Sobol's model state.
      gr._model_key = "GPEI"
      gr._model_kwargs = gpei._model_kwargs
      gr._bridge_kwargs = gpei._bridge_kwargs
      gr._model_state_after_gen = {}
      gpei_restored = get_model_from_generator_run(
          gr, experiment=exp, data=get_branin_data(), models_enum=Models
      )
      for key in gpei.__dict__:
          self.assertIn(key, gpei_restored.__dict__)
          original, restored = (
              gpei.__dict__[key],
              gpei_restored.__dict__[key],
          )
          # Fit times are set in instantiation so not same and model compared below.
          if key in [
              "fit_time",
              "fit_time_since_gen",
              "model",
              "training_data",
          ]:
              continue  # Fit times are set in instantiation so won't be same.
          if isinstance(original, OrderedDict) and isinstance(
              restored, OrderedDict
          ):
              original, restored = list(original.keys()), list(
                  restored.keys()
              )
          if isinstance(original, Model) and isinstance(
              restored, Model
          ):
              continue  # Model equality is tough to compare.
              self.assertEqual(original, restored)
  
      for key in gpei.model.__dict__:
          self.assertIn(key, gpei_restored.model.__dict__)
          original, restored = (
              gpei.model.__dict__[key],
              gpei_restored.model.__dict__[key],
          )
          # Botorch model equality is tough to compare and training data
          # is unnecessary to compare, because data passed to model was the same
          if key in ["_model", "warm_start_refitting", "Xs", "Ys"]:
              continue
          self.assertEqual(original, restored)
  
imports: from collections import OrderedDict
  import numpy
  import pandas
  import torch
  from ax.core.data import Data
  from ax.core.observation import ObservationFeatures
  from ax.modelbridge.discrete import DiscreteModelBridge
  from ax.modelbridge.random import RandomModelBridge
  from ax.modelbridge.registry import Cont_X_trans, get_model_from_generator_run, MODEL_KEY_TO_MODEL_SETUP, Models, Y_trans
  from ax.modelbridge.torch import TorchModelBridge
  from ax.models.base import Model
  from ax.models.discrete.eb_thompson import EmpiricalBayesThompsonSampler
  from ax.models.discrete.thompson import ThompsonSampler
  from ax.models.random.alebo_initializer import ALEBOInitializer
  from ax.models.torch.alebo import ALEBO
  from ax.models.torch.botorch_modular.acquisition import Acquisition
  from ax.models.torch.botorch_modular.kernels import ScaleMaternKernel
  from ax.models.torch.botorch_modular.model import BoTorchModel, SurrogateSpec
  from ax.models.torch.botorch_modular.surrogate import Surrogate
  from ax.models.torch.botorch_moo import MultiObjectiveBotorchModel
  from ax.utils.common.constants import Keys
  from ax.utils.common.kwargs import get_function_argument_names
  from ax.utils.common.testutils import TestCase
  from ax.utils.testing.core_stubs import get_branin_data, get_branin_experiment, get_branin_experiment_with_status_quo_trials, get_branin_optimization_config, get_factorial_experiment
  from ax.utils.testing.mock import fast_botorch_optimize
  from botorch.acquisition.monte_carlo import qExpectedImprovement
  from botorch.models.fully_bayesian import SaasFullyBayesianSingleTaskGP
  from botorch.models.fully_bayesian_multitask import SaasFullyBayesianMultiTaskGP
  from botorch.models.gp_regression import SingleTaskGP
  from botorch.models.model_list_gp_regression import ModelListGP
  from botorch.models.multitask import MultiTaskGP
  from botorch.utils.types import DEFAULT
  from gpytorch.kernels.matern_kernel import MaternKernel
  from gpytorch.kernels.scale_kernel import ScaleKernel
  from gpytorch.mlls.exact_marginal_log_likelihood import ExactMarginalLogLikelihood
  from gpytorch.priors.torch_priors import GammaPrior
label_desc: Use image node_id 10 for calling the ModelRegistryTest obj's underlying member method code with example usage: obj.test_get_model_from_generator_run() without return types