class_docstr: Base class for Tasks that are generated using the Profet algorithm.
  
  For more information on Profet, see original paper at https://arxiv.org/abs/1905.12982.
  
  Klein, Aaron, Zhenwen Dai, Frank Hutter, Neil Lawrence, and Javier Gonzalez. "Meta-surrogate benchmarking for
  hyperparameter optimization." Advances in Neural Information Processing Systems 32 (2019): 6270-6280.
  
  Parameters
  ----------
  max_trials : int, optional
      Max number of trials to run, by default 100
  input_dir : Union[Path, str], optional
      Input directory containing the data used to train the meta-model, by default None.
  checkpoint_dir : Union[Path, str], optional
      Directory used to save/load trained meta-models, by default None.
  model_config : MetaModelConfig, optional
      Configuration options for the training of the meta-model, by default None
  device : str, optional
      The device to use for training, by default None.
  with_grad : bool, optional
      Whether the task should also return the gradients of the objective function with respect to
      the inputs. Defaults to `False`.
class_docstr_tok: ["Base","class","for","Tasks","that","are","generated","using","the","Profet","algorithm",".","For","more","information","on","Profet",",","see","original","paper","at","https",":","\/\/arxiv.org\/abs\/1905.12982",".","Klein",",","Aaron",",","Zhenwen","Dai",",","Frank","Hutter",",","Neil","Lawrence",",","and","Javier","Gonzalez",".","``","Meta-surrogate","benchmarking","for","hyperparameter","optimization",".","''","Advances","in","Neural","Information","Processing","Systems","32","(","2019",")",":","6270-6280",".","Parameters","--","--","--","--","--","max_trials",":","int",",","optional","Max","number","of","trials","to","run",",","by","default","100","input_dir",":","Union","[","Path",",","str","]",",","optional","Input","directory","containing","the","data","used","to","train","the","meta-model",",","by","default","None",".","checkpoint_dir",":","Union","[","Path",",","str","]",",","optional","Directory","used","to","save\/load","trained","meta-models",",","by","default","None",".","model_config",":","MetaModelConfig",",","optional","Configuration","options","for","the","training","of","the","meta-model",",","by","default","None","device",":","str",",","optional","The","device","to","use","for","training",",","by","default","None",".","with_grad",":","bool",",","optional","Whether","the","task","should","also","return","the","gradients","of","the","objective","function","with","respect","to","the","inputs",".","Defaults","to","`","False","`","."]
code: def __init__(
      self,
      max_trials: int = 100,
      input_dir: Union[Path, str] = "profet_data",
      checkpoint_dir: Union[Path, str] = None,
      model_config: MetaModelConfig = None,
      device: Union[str, Any] = None,
      with_grad: bool = False,
  ):
      super().__init__(max_trials=max_trials)
      self.input_dir = Path(input_dir)
      self.checkpoint_dir = Path(
          checkpoint_dir or self.input_dir / "checkpoints"
      )
  
      # The config for the training of the meta-model.
      # NOTE: the train config is used to determine the hash of the task.
      if model_config is None:
          # NOTE: This type error is safe to ignore: the benchmark argument will have been set in
          # each ModelConfig subclass.
          self.model_config = self.ModelConfig()  # type: ignore
      elif isinstance(model_config, dict):
          self.model_config = self.ModelConfig(**model_config)
      elif not isinstance(model_config, self.ModelConfig):
          # If passed a model config, for example through deserializing the configuration,
          # then convert it back to the right type, so the class attributes are correct.
          self.model_config = self.ModelConfig(**asdict(model_config))
      else:
          self.model_config = model_config
  
      assert isinstance(self.model_config, self.ModelConfig)
  
      self.seed = self.model_config.seed
  
      self.with_grad = with_grad
      # The parameters that have an influence over the training of the meta-model are used to
      # create the filename where the model will be saved.
      task_hash_params = asdict(self.model_config)
      logger.info(f"Task hash params: {task_hash_params}")
      task_hash = compute_identity(**task_hash_params)
  
      filename = f"{task_hash}.pkl"
  
      self.checkpoint_file = self.checkpoint_dir / filename
      logger.info(
          f"Checkpoint file for this task: {self.checkpoint_file}"
      )
  
      if isinstance(device, torch.device):
          self.device = device
      else:
          self.device = torch.device(
              device or ("cuda" if torch.cuda.is_available() else "cpu")
          )
  
      # NOTE: Need to control the randomness that's happening inside *both* the training
      # function, as well as the loading function (since `load_task_network`` instantiates a model
      # and then loads the weights, it also affects the global rng state of pytorch).
  
      with make_reproducible(self.seed):
          if os.path.exists(self.checkpoint_file):
              logger.info(
                  f"Model has already been trained: loading it from file {self.checkpoint_file}."
              )
              self.net, h = self.model_config.load_task_network(
                  self.checkpoint_file
              )
          else:
              warnings.warn(
                  RuntimeWarning(
                      f"Checkpoint file {self.checkpoint_file} doesn't exist: re-training the "
                      f"model. (This may take a *very* long time!)"
                  )
              )
              logger.info(f"Task hash params: {task_hash_params}")
              self.checkpoint_file.parent.mkdir(
                  exist_ok=True, parents=True
              )
  
              # Need to re-train the meta-model and sample this task.
              self.net, h = self.model_config.get_task_network(
                  self.input_dir
              )
  
      # Numpy random state. Currently only used in `sample()`
      self._np_rng_state = np.random.RandomState(self.seed)
  
      self.h: np.ndarray = np.array(h)
      self.model_config.save_task_network(
          self.checkpoint_file, self.net, self.h
      )
  
      self.net = self.net.to(device=self.device, dtype=torch.float32)
      self.net.eval()
  
      self.h_tensor = torch.as_tensor(
          self.h, dtype=torch.float32, device=self.device
      )
  
      self._space: Optional[Space] = None
      self.name = f"profet.{type(self).__qualname__.lower()}_{self.model_config.task_id}"
      self.transformed_space = transformer.build_required_space(
          self.space,
          type_requirement="real",
          shape_requirement="flattened",
          dist_requirement="linear",
      )
  
imports: import os
  import random
  import warnings
  from abc import ABC
  from contextlib import contextmanager
  from dataclasses import asdict
  from logging import getLogger
  from pathlib import Path
  from typing import Any, ClassVar, Dict, List, Optional, Type, Union
  import numpy
  from orion.algo.space import Space
  from orion.benchmark.task.base import BenchmarkTask
  from orion.benchmark.task.profet.model_utils import MetaModelConfig
  from orion.core.io.space_builder import SpaceBuilder
  from orion.core.utils import compute_identity
  from orion.core.utils.flatten import flatten
  from orion.core.utils.format_trials import dict_to_trial
  from orion.core.worker import transformer
label_desc: Use image node_id 1 to create a new ProfetTask object from inherited base classes: BenchmarkTask, ABC with example: obj = ProfetTask(max_trials, input_dir, checkpoint_dir, model_config, device, with_grad)