code: def test_index_to_mask():
      assert str(IndexToMask()) == (
          "IndexToMask(attrs=None, sizes=None, " "replace=False)"
      )
  
      edge_index = torch.tensor(
          [[0, 1, 1, 2, 2, 3, 3, 4], [1, 0, 2, 1, 3, 2, 4, 3]]
      )
      train_index = torch.arange(0, 3)
      test_index = torch.arange(3, 5)
      data = Data(
          edge_index=edge_index,
          train_index=train_index,
          test_index=test_index,
          num_nodes=5,
      )
  
      out = IndexToMask(replace=True)(data)
      assert len(out) == len(data)
      assert out.train_mask.tolist() == [True, True, True, False, False]
      assert out.test_mask.tolist() == [False, False, False, True, True]
  
      out = IndexToMask(replace=False)(data)
      assert len(out) == len(data) + 2
  
      out = IndexToMask(sizes=6, replace=True)(data)
      assert out.train_mask.tolist() == [
          True,
          True,
          True,
          False,
          False,
          False,
      ]
      assert out.test_mask.tolist() == [
          False,
          False,
          False,
          True,
          True,
          False,
      ]
  
      out = IndexToMask(attrs="train_index")(data)
      assert len(out) == len(data) + 1
      assert "train_index" in out
      assert "train_mask" in out
      assert "test_index" in out
      assert "test_mask" not in out
  
imports: import torch
  from torch_geometric.data import Data, HeteroData
  from torch_geometric.transforms import IndexToMask, MaskToIndex
label_desc: Use image node_id 1 for calling a global function with example usage: test_index_to_mask() without return types