code: def test_multi_task(self):
      dataset_1 = make_dataset(outcome_names=["y"])
      dataset_2 = make_dataset(outcome_names=["z"])
      dataset_3 = make_dataset(has_yvar=True, outcome_names=["z"])
      dataset_4 = make_dataset(has_yvar=True, outcome_names=["y"])
      # Test validation.
      with self.assertRaisesRegex(
          UnsupportedError, "containing more than one outcome"
      ):
          MultiTaskDataset(
              datasets=[make_dataset(m=2)], target_outcome_name="y0"
          )
      with self.assertRaisesRegex(
          UnsupportedError, "multiple datasets for the same outcome"
      ):
          MultiTaskDataset(
              datasets=[dataset_1, dataset_1], target_outcome_name="y"
          )
      with self.assertRaisesRegex(
          InputDataError, "Target outcome is not present"
      ):
          MultiTaskDataset(
              datasets=[dataset_1], target_outcome_name="z"
          )
      with self.assertRaisesRegex(
          UnsupportedError, "modeling batched inputs"
      ):
          MultiTaskDataset(
              datasets=[make_dataset(batch_shape=torch.Size([2]))],
              target_outcome_name="y0",
          )
      with self.assertRaisesRegex(
          InputDataError, "names of the task features"
      ):
          MultiTaskDataset(
              datasets=[
                  dataset_1,
                  make_dataset(
                      feature_names=["x1", "x3"], outcome_names=["z"]
                  ),
              ],
              target_outcome_name="z",
              task_feature_index=1,
          )
      with self.assertRaisesRegex(
          UnsupportedError,
          "all or none of the datasets to have a Yvar.",
      ):
          MultiTaskDataset(
              datasets=[dataset_1, dataset_3], target_outcome_name="z"
          )
  
      # Test correct construction.
      mt_dataset = MultiTaskDataset(
          datasets=[dataset_1, dataset_2],
          target_outcome_name="z",
      )
      self.assertEqual(len(mt_dataset.datasets), 2)
      self.assertIsNone(mt_dataset.task_feature_index)
      self.assertIs(mt_dataset.datasets["y"], dataset_1)
      self.assertIs(mt_dataset.datasets["z"], dataset_2)
      self.assertIsNone(mt_dataset.Yvar)
      self.assertFalse(mt_dataset.has_heterogeneous_features)
      expected_X = torch.cat(
          [
              torch.cat([dataset_1.X, torch.ones(3, 1)], dim=-1),
              torch.cat([dataset_2.X, torch.zeros(3, 1)], dim=-1),
          ],
          dim=0,
      )
      expected_Y = torch.cat(
          [ds.Y for ds in [dataset_1, dataset_2]], dim=0
      )
      self.assertTrue(torch.equal(expected_X, mt_dataset.X))
      self.assertTrue(torch.equal(expected_Y, mt_dataset.Y))
      self.assertIs(
          mt_dataset.get_dataset_without_task_feature(outcome_name="y"),
          dataset_1,
      )
  
      # Test with Yvar and target_feature_index.
      mt_dataset = MultiTaskDataset(
          datasets=[dataset_3, dataset_4],
          target_outcome_name="z",
          task_feature_index=1,
      )
      self.assertEqual(mt_dataset.task_feature_index, 1)
      expected_X_2 = torch.cat([dataset_3.X, dataset_4.X], dim=0)
      expected_Yvar_2 = torch.cat(
          [dataset_3.Yvar, dataset_4.Yvar], dim=0
      )
      self.assertTrue(torch.equal(expected_X_2, mt_dataset.X))
      self.assertTrue(torch.equal(expected_Yvar_2, mt_dataset.Yvar))
      # Check that the task feature is removed correctly.
      ds_3_no_task = mt_dataset.get_dataset_without_task_feature(
          outcome_name="z"
      )
      self.assertTrue(torch.equal(ds_3_no_task.X, dataset_3.X[:, :1]))
      self.assertTrue(torch.equal(ds_3_no_task.Y, dataset_3.Y))
      self.assertTrue(torch.equal(ds_3_no_task.Yvar, dataset_3.Yvar))
      self.assertEqual(
          ds_3_no_task.feature_names, dataset_3.feature_names[:1]
      )
      self.assertEqual(
          ds_3_no_task.outcome_names, dataset_3.outcome_names
      )
  
      # Test from_joint_dataset.
      sort_idcs = [
          3,
          4,
          5,
          0,
          1,
          2,
      ]  # X & Y will get sorted based on task feature.
      for outcome_names_per_task in [None, {0: "x", 1: "y"}]:
          joint_dataset = SupervisedDataset(
              X=expected_X,
              Y=expected_Y,
              feature_names=["x0", "x1", "task"],
              outcome_names=["z"],
          )
          mt_dataset = MultiTaskDataset.from_joint_dataset(
              dataset=joint_dataset,
              task_feature_index=-1,
              target_task_value=0,
              outcome_names_per_task=outcome_names_per_task,
          )
          self.assertEqual(len(mt_dataset.datasets), 2)
          if outcome_names_per_task is None:
              self.assertEqual(
                  list(mt_dataset.datasets.keys()), ["z", "task_1"]
              )
              self.assertEqual(mt_dataset.target_outcome_name, "z")
          else:
              self.assertEqual(
                  list(mt_dataset.datasets.keys()), ["x", "y"]
              )
              self.assertEqual(mt_dataset.target_outcome_name, "x")
  
          self.assertTrue(
              torch.equal(mt_dataset.X, expected_X[sort_idcs])
          )
          self.assertTrue(
              torch.equal(mt_dataset.Y, expected_Y[sort_idcs])
          )
          self.assertTrue(
              torch.equal(
                  mt_dataset.datasets[mt_dataset.target_outcome_name].Y,
                  dataset_2.Y,
              )
          )
          self.assertIsNone(mt_dataset.Yvar)
      with self.assertRaisesRegex(
          UnsupportedError, "more than one outcome"
      ):
          MultiTaskDataset.from_joint_dataset(
              dataset=make_dataset(m=2),
              task_feature_index=-1,
              target_task_value=0,
          )
  
      # With heterogeneous feature sets.
      dataset_5 = make_dataset(d=1, outcome_names=["z"])
      mt_dataset = MultiTaskDataset(
          datasets=[dataset_1, dataset_5],
          target_outcome_name="y",
      )
      self.assertTrue(mt_dataset.has_heterogeneous_features)
      with self.assertRaisesRegex(
          UnsupportedError, "datasets with heterogeneous feature sets"
      ):
          mt_dataset.X
  
file_path: botorch/test/utils/test_datasets.py
imports: from typing import List, Optional
  import torch
  from botorch.exceptions.errors import InputDataError, UnsupportedError
  from botorch.utils.containers import DenseContainer, SliceContainer
  from botorch.utils.datasets import ContextualDataset, FixedNoiseDataset, MultiTaskDataset, RankingDataset, SupervisedDataset
  from botorch.utils.testing import BotorchTestCase
  from torch import rand, randperm, Size, stack, Tensor, tensor
label_desc: Use image node_id 4 for calling the TestDatasets obj's underlying member method code with example usage: obj.test_multi_task() without return types