class_docstr: A model bridge for using torch-based models that fit on MapData. Most
  of the `TorchModelBridge` functionality is retained, except that this
  class should be used in the case where `model` makes use of map_key values.
  For example, the use case of fitting a joint surrogate model on
  `(parameters, map_key)`, while candidate generation is only for `parameters`.
class_docstr_tok: ["A","model","bridge","for","using","torch-based","models","that","fit","on","MapData",".","Most","of","the","`","TorchModelBridge","`","functionality","is","retained",",","except","that","this","class","should","be","used","in","the","case","where","`","model","`","makes","use","of","map_key","values",".","For","example",",","the","use","case","of","fitting","a","joint","surrogate","model","on","`","(","parameters",",","map_key",")","`",",","while","candidate","generation","is","only","for","`","parameters","`","."]
docstr: Applies transforms and fits model.
  
  Args:
      experiment: Is used to get arm parameters. Is not mutated.
      search_space: Search space for fitting the model. Constraints need
          not be the same ones used in gen.
      data: Ax Data.
      model: Interface will be specified in subclass. If model requires
          initialization, that should be done prior to its use here.
      transforms: List of uninitialized transform classes. Forward
          transforms will be applied in this order, and untransforms in
          the reverse order.
      transform_configs: A dictionary from transform name to the
          transform config dictionary.
      torch_dtype: Torch data type.
      torch_device: Torch device.
      status_quo_name: Name of the status quo arm. Can only be used if
          Data has a single set of ObservationFeatures corresponding to
          that arm.
      status_quo_features: ObservationFeatures to use as status quo.
          Either this or status_quo_name should be specified, not both.
      optimization_config: Optimization config defining how to optimize
          the model.
      fit_out_of_design: If specified, all training data is returned.
          Otherwise, only in design points are returned.
      fit_on_init: Whether to fit the model on initialization. This can
          be used to skip model fitting when a fitted model is not needed.
          To fit the model afterwards, use `_process_and_transform_data`
          to get the transformed inputs and call `_fit_if_implemented` with
          the transformed inputs.
      default_model_gen_options: Options passed down to `model.gen(...)`.
      map_data_limit_rows_per_metric: Subsample the map data so that the
          total number of rows per metric is limited by this value.
      map_data_limit_rows_per_group: Subsample the map data so that the
          number of rows in the `map_key` column for each (arm, metric)
          is limited by this value.
docstr_tok: ["Applies","transforms","and","fits","model",".","Args",":","experiment",":","Is","used","to","get","arm","parameters",".","Is","not","mutated",".","search_space",":","Search","space","for","fitting","the","model",".","Constraints","need","not","be","the","same","ones","used","in","gen.","data",":","Ax","Data",".","model",":","Interface","will","be","specified","in","subclass",".","If","model","requires","initialization",",","that","should","be","done","prior","to","its","use","here",".","transforms",":","List","of","uninitialized","transform","classes",".","Forward","transforms","will","be","applied","in","this","order",",","and","untransforms","in","the","reverse","order",".","transform_configs",":","A","dictionary","from","transform","name","to","the","transform","config","dictionary",".","torch_dtype",":","Torch","data","type",".","torch_device",":","Torch","device",".","status_quo_name",":","Name","of","the","status","quo","arm",".","Can","only","be","used","if","Data","has","a","single","set","of","ObservationFeatures","corresponding","to","that","arm",".","status_quo_features",":","ObservationFeatures","to","use","as","status","quo",".","Either","this","or","status_quo_name","should","be","specified",",","not","both",".","optimization_config",":","Optimization","config","defining","how","to","optimize","the","model",".","fit_out_of_design",":","If","specified",",","all","training","data","is","returned",".","Otherwise",",","only","in","design","points","are","returned",".","fit_on_init",":","Whether","to","fit","the","model","on","initialization",".","This","can","be","used","to","skip","model","fitting","when","a","fitted","model","is","not","needed",".","To","fit","the","model","afterwards",",","use","`","_process_and_transform_data","`","to","get","the","transformed","inputs","and","call","`","_fit_if_implemented","`","with","the","transformed","inputs",".","default_model_gen_options",":","Options","passed","down","to","`","model.gen","(","...",")","`",".","map_data_limit_rows_per_metric",":","Subsample","the","map","data","so","that","the","total","number","of","rows","per","metric","is","limited","by","this","value",".","map_data_limit_rows_per_group",":","Subsample","the","map","data","so","that","the","number","of","rows","in","the","`","map_key","`","column","for","each","(","arm",",","metric",")","is","limited","by","this","value","."]
code: def __init__(
      self,
      experiment: Experiment,
      search_space: SearchSpace,
      data: Data,
      model: TorchModel,
      transforms: List[Type[Transform]],
      transform_configs: Optional[Dict[str, TConfig]] = None,
      torch_dtype: Optional[torch.dtype] = None,
      torch_device: Optional[torch.device] = None,
      status_quo_name: Optional[str] = None,
      status_quo_features: Optional[ObservationFeatures] = None,
      optimization_config: Optional[OptimizationConfig] = None,
      fit_out_of_design: bool = False,
      fit_on_init: bool = True,
      default_model_gen_options: Optional[TConfig] = None,
      map_data_limit_rows_per_metric: Optional[int] = None,
      map_data_limit_rows_per_group: Optional[int] = None,
  ) -> None:
      """
      Applies transforms and fits model.
  
      Args:
          experiment: Is used to get arm parameters. Is not mutated.
          search_space: Search space for fitting the model. Constraints need
              not be the same ones used in gen.
          data: Ax Data.
          model: Interface will be specified in subclass. If model requires
              initialization, that should be done prior to its use here.
          transforms: List of uninitialized transform classes. Forward
              transforms will be applied in this order, and untransforms in
              the reverse order.
          transform_configs: A dictionary from transform name to the
              transform config dictionary.
          torch_dtype: Torch data type.
          torch_device: Torch device.
          status_quo_name: Name of the status quo arm. Can only be used if
              Data has a single set of ObservationFeatures corresponding to
              that arm.
          status_quo_features: ObservationFeatures to use as status quo.
              Either this or status_quo_name should be specified, not both.
          optimization_config: Optimization config defining how to optimize
              the model.
          fit_out_of_design: If specified, all training data is returned.
              Otherwise, only in design points are returned.
          fit_on_init: Whether to fit the model on initialization. This can
              be used to skip model fitting when a fitted model is not needed.
              To fit the model afterwards, use `_process_and_transform_data`
              to get the transformed inputs and call `_fit_if_implemented` with
              the transformed inputs.
          default_model_gen_options: Options passed down to `model.gen(...)`.
          map_data_limit_rows_per_metric: Subsample the map data so that the
              total number of rows per metric is limited by this value.
          map_data_limit_rows_per_group: Subsample the map data so that the
              number of rows in the `map_key` column for each (arm, metric)
              is limited by this value.
      """
  
      if not isinstance(data, MapData):
          raise ValueError(
              "`MapTorchModelBridge expects `MapData` instead of `Data`."
          )
      # pyre-fixme[4]: Attribute must be annotated.
      self._map_key_features = data.map_keys
      self._map_data_limit_rows_per_metric = (
          map_data_limit_rows_per_metric
      )
      self._map_data_limit_rows_per_group = (
          map_data_limit_rows_per_group
      )
  
      super().__init__(
          experiment=experiment,
          search_space=search_space,
          data=data,
          model=model,
          transforms=transforms,
          transform_configs=transform_configs,
          torch_dtype=torch_dtype,
          torch_device=torch_device,
          status_quo_name=status_quo_name,
          status_quo_features=status_quo_features,
          optimization_config=optimization_config,
          fit_out_of_design=fit_out_of_design,
          fit_on_init=fit_on_init,
          default_model_gen_options=default_model_gen_options,
      )
  
imports: from typing import Any, Dict, List, Optional, Tuple, Type
  import numpy
  import torch
  from ax.core.data import Data
  from ax.core.experiment import Experiment
  from ax.core.map_data import MapData
  from ax.core.observation import Observation, ObservationData, ObservationFeatures, observations_from_map_data, separate_observations
  from ax.core.optimization_config import OptimizationConfig
  from ax.core.search_space import SearchSpace
  from ax.core.types import TCandidateMetadata
  from ax.modelbridge.base import GenResults
  from ax.modelbridge.modelbridge_utils import array_to_observation_data, observation_features_to_array, parse_observation_features
  from ax.modelbridge.torch import FIT_MODEL_ERROR, TorchModelBridge
  from ax.modelbridge.transforms.base import Transform
  from ax.models.torch_base import TorchModel
  from ax.models.types import TConfig
  from ax.utils.common.constants import Keys
  from ax.utils.common.typeutils import not_none
label_desc: Use image node_id 1 to create a new MapTorchModelBridge object from inherited base classes: TorchModelBridge with example: obj = MapTorchModelBridge(experiment, search_space, data, model, transforms, transform_configs, torch_dtype, torch_device, status_quo_name, status_quo_features, optimization_config, fit_out_of_design, fit_on_init, default_model_gen_options, map_data_limit_rows_per_metric, map_data_limit_rows_per_group)