class_docstr: A model bridge for using torch-based models.
  
  Specifies an interface that is implemented by TorchModel. In particular,
  model should have methods fit, predict, and gen. See TorchModel for the
  API for each of these methods.
  
  Requires that all parameters have been transformed to RangeParameters
  or FixedParameters with float type and no log scale.
  
  This class converts Ax parameter types to torch tensors before passing
  them to the model.
class_docstr_tok: ["A","model","bridge","for","using","torch-based","models",".","Specifies","an","interface","that","is","implemented","by","TorchModel",".","In","particular",",","model","should","have","methods","fit",",","predict",",","and","gen.","See","TorchModel","for","the","API","for","each","of","these","methods",".","Requires","that","all","parameters","have","been","transformed","to","RangeParameters","or","FixedParameters","with","float","type","and","no","log","scale",".","This","class","converts","Ax","parameter","types","to","torch","tensors","before","passing","them","to","the","model","."]
docstr: Helper for consolidating some common argument processing between
  fit and cross validate methods. Extracts datasets and candidate metadate
  from observations and the search space digest from the search space.
  
  Args:
      search_space: A transformed search space for fitting the model.
      observations: The observations to fit the model with. These should
          also be transformed.
      parameters: Names of parameters to be used in the model. Defaults to
          all parameters in the search space.
      update_outcomes_and_parameters: Whether to update `self.outcomes` with
          all outcomes found in the observations and `self.parameters` with
          all parameters in the search space. Typically only used in `_fit`.
  
  Returns:
      The datasets & metadata extracted from the observations and the
      search space digest.
docstr_tok: ["Helper","for","consolidating","some","common","argument","processing","between","fit","and","cross","validate","methods",".","Extracts","datasets","and","candidate","metadate","from","observations","and","the","search","space","digest","from","the","search","space",".","Args",":","search_space",":","A","transformed","search","space","for","fitting","the","model",".","observations",":","The","observations","to","fit","the","model","with",".","These","should","also","be","transformed",".","parameters",":","Names","of","parameters","to","be","used","in","the","model",".","Defaults","to","all","parameters","in","the","search","space",".","update_outcomes_and_parameters",":","Whether","to","update","`","self.outcomes","`","with","all","outcomes","found","in","the","observations","and","`","self.parameters","`","with","all","parameters","in","the","search","space",".","Typically","only","used","in","`","_fit","`",".","Returns",":","The","datasets","&","metadata","extracted","from","the","observations","and","the","search","space","digest","."]
code: def _get_fit_args(
      self,
      search_space: SearchSpace,
      observations: List[Observation],
      parameters: Optional[List[str]],
      update_outcomes_and_parameters: bool,
  ) -> Tuple[
      List[SupervisedDataset],
      Optional[List[List[TCandidateMetadata]]],
      SearchSpaceDigest,
  ]:
      """Helper for consolidating some common argument processing between
      fit and cross validate methods. Extracts datasets and candidate metadate
      from observations and the search space digest from the search space.
  
      Args:
          search_space: A transformed search space for fitting the model.
          observations: The observations to fit the model with. These should
              also be transformed.
          parameters: Names of parameters to be used in the model. Defaults to
              all parameters in the search space.
          update_outcomes_and_parameters: Whether to update `self.outcomes` with
              all outcomes found in the observations and `self.parameters` with
              all parameters in the search space. Typically only used in `_fit`.
  
      Returns:
          The datasets & metadata extracted from the observations and the
          search space digest.
      """
      self._last_observations = observations
      if update_outcomes_and_parameters:
          self.parameters = list(search_space.parameters.keys())
      if parameters is None:
          parameters = self.parameters
      all_metric_names: Set[str] = set()
      observation_features, observation_data = separate_observations(
          observations
      )
      # Only update outcomes if fitting a model on tracking metrics. Otherwise,
      # we will only fit models to the outcomes that are extracted from optimization
      # config in ModelBridge.__init__.
      if update_outcomes_and_parameters and self._fit_tracking_metrics:
          for od in observation_data:
              all_metric_names.update(od.metric_names)
          self.outcomes = sorted(
              all_metric_names
          )  # Deterministic order
      # Get all relevant information on the parameters
      search_space_digest = extract_search_space_digest(
          search_space=search_space, param_names=self.parameters
      )
      # Convert observations to datasets
      (
          datasets,
          ordered_outcomes,
          candidate_metadata,
      ) = self._convert_observations(
          observation_data=observation_data,
          observation_features=observation_features,
          outcomes=self.outcomes,
          parameters=parameters,
          search_space_digest=search_space_digest,
      )
      if update_outcomes_and_parameters:
          self.outcomes = ordered_outcomes
      else:
          assert (
              ordered_outcomes == self.outcomes
          ), f"Unexpected ordering of outcomes: {ordered_outcomes} != {self.outcomes}"
      return datasets, candidate_metadata, search_space_digest
  
imports: from __future__ import annotations
  from collections import defaultdict
  from copy import deepcopy
  from logging import Logger
  from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Type, Union
  import numpy
  import torch
  from ax.core.arm import Arm
  from ax.core.data import Data
  from ax.core.experiment import Experiment
  from ax.core.generator_run import extract_arm_predictions
  from ax.core.metric import Metric
  from ax.core.observation import Observation, ObservationData, ObservationFeatures, separate_observations
  from ax.core.optimization_config import MultiObjectiveOptimizationConfig, OptimizationConfig
  from ax.core.outcome_constraint import ComparisonOp, ObjectiveThreshold, ScalarizedOutcomeConstraint
  from ax.core.search_space import SearchSpace
  from ax.core.types import TCandidateMetadata, TModelPredictArm
  from ax.exceptions.core import DataRequiredError, UnsupportedError
  from ax.modelbridge.base import gen_arms, GenResults, ModelBridge
  from ax.modelbridge.modelbridge_utils import array_to_observation_data, extract_objective_thresholds, extract_objective_weights, extract_outcome_constraints, extract_parameter_constraints, extract_risk_measure, extract_search_space_digest, get_fixed_features, observation_data_to_array, observation_features_to_array, parse_observation_features, pending_observations_as_array_list, process_contextual_datasets, SearchSpaceDigest, transform_callback, validate_and_apply_final_transform
  from ax.modelbridge.transforms.base import Transform
  from ax.modelbridge.transforms.cast import Cast
  from ax.models.torch.botorch_modular.model import BoTorchModel
  from ax.models.torch.botorch_moo import MultiObjectiveBotorchModel
  from ax.models.torch.botorch_moo_defaults import infer_objective_thresholds
  from ax.models.torch_base import TorchModel, TorchOptConfig
  from ax.models.types import TConfig
  from ax.utils.common.logger import get_logger
  from ax.utils.common.typeutils import not_none
  from botorch.utils.datasets import MultiTaskDataset, SupervisedDataset
  from torch import Tensor
label_desc: Use image node_id 12 for calling the TorchModelBridge obj's underlying member method code with example usage: obj._get_fit_args(search_space, observations, parameters, update_outcomes_and_parameters) and returns: datasets,  candidate_metadata,  search_space_digest