code: def test_random_node_split(num_splits):
      num_nodes, num_classes = 1000, 4
      x = torch.randn(num_nodes, 16)
      y = torch.randint(num_classes, (num_nodes,), dtype=torch.long)
      data = Data(x=x, y=y)
  
      transform = RandomNodeSplit(
          split="train_rest",
          num_splits=num_splits,
          num_val=100,
          num_test=200,
      )
      assert str(transform) == "RandomNodeSplit(split=train_rest)"
      data = transform(data)
      assert len(data) == 5
  
      train_mask = data.train_mask
      train_mask = (
          train_mask.unsqueeze(-1) if num_splits == 1 else train_mask
      )
      assert train_mask.size() == (num_nodes, num_splits)
      val_mask = data.val_mask
      val_mask = val_mask.unsqueeze(-1) if num_splits == 1 else val_mask
      assert val_mask.size() == (num_nodes, num_splits)
      test_mask = data.test_mask
      test_mask = (
          test_mask.unsqueeze(-1) if num_splits == 1 else test_mask
      )
      assert test_mask.size() == (num_nodes, num_splits)
  
      for i in range(train_mask.size(-1)):
          assert train_mask[:, i].sum() == num_nodes - 100 - 200
          assert val_mask[:, i].sum() == 100
          assert test_mask[:, i].sum() == 200
          assert (
              train_mask[:, i] & val_mask[:, i] & test_mask[:, i]
          ).sum() == 0
          assert (
              train_mask[:, i] | val_mask[:, i] | test_mask[:, i]
          ).sum() == num_nodes
  
      transform = RandomNodeSplit(
          split="train_rest",
          num_splits=num_splits,
          num_val=0.1,
          num_test=0.2,
      )
      data = transform(data)
  
      train_mask = data.train_mask
      train_mask = (
          train_mask.unsqueeze(-1) if num_splits == 1 else train_mask
      )
      val_mask = data.val_mask
      val_mask = val_mask.unsqueeze(-1) if num_splits == 1 else val_mask
      test_mask = data.test_mask
      test_mask = (
          test_mask.unsqueeze(-1) if num_splits == 1 else test_mask
      )
  
      for i in range(train_mask.size(-1)):
          assert train_mask[:, i].sum() == num_nodes - 100 - 200
          assert val_mask[:, i].sum() == 100
          assert test_mask[:, i].sum() == 200
          assert (
              train_mask[:, i] & val_mask[:, i] & test_mask[:, i]
          ).sum() == 0
          assert (
              train_mask[:, i] | val_mask[:, i] | test_mask[:, i]
          ).sum() == num_nodes
  
      transform = RandomNodeSplit(
          split="test_rest",
          num_splits=num_splits,
          num_train_per_class=10,
          num_val=100,
      )
      assert str(transform) == "RandomNodeSplit(split=test_rest)"
      data = transform(data)
      assert len(data) == 5
  
      train_mask = data.train_mask
      train_mask = (
          train_mask.unsqueeze(-1) if num_splits == 1 else train_mask
      )
      val_mask = data.val_mask
      val_mask = val_mask.unsqueeze(-1) if num_splits == 1 else val_mask
      test_mask = data.test_mask
      test_mask = (
          test_mask.unsqueeze(-1) if num_splits == 1 else test_mask
      )
  
      for i in range(train_mask.size(-1)):
          assert train_mask[:, i].sum() == 10 * num_classes
          assert val_mask[:, i].sum() == 100
          assert (
              test_mask[:, i].sum()
              == num_nodes - 10 * num_classes - 100
          )
          assert (
              train_mask[:, i] & val_mask[:, i] & test_mask[:, i]
          ).sum() == 0
          assert (
              train_mask[:, i] | val_mask[:, i] | test_mask[:, i]
          ).sum() == num_nodes
  
      transform = RandomNodeSplit(
          split="test_rest",
          num_splits=num_splits,
          num_train_per_class=10,
          num_val=0.1,
      )
      data = transform(data)
  
      train_mask = data.train_mask
      train_mask = (
          train_mask.unsqueeze(-1) if num_splits == 1 else train_mask
      )
      val_mask = data.val_mask
      val_mask = val_mask.unsqueeze(-1) if num_splits == 1 else val_mask
      test_mask = data.test_mask
      test_mask = (
          test_mask.unsqueeze(-1) if num_splits == 1 else test_mask
      )
  
      for i in range(train_mask.size(-1)):
          assert train_mask[:, i].sum() == 10 * num_classes
          assert val_mask[:, i].sum() == 100
          assert (
              test_mask[:, i].sum()
              == num_nodes - 10 * num_classes - 100
          )
          assert (
              train_mask[:, i] & val_mask[:, i] & test_mask[:, i]
          ).sum() == 0
          assert (
              train_mask[:, i] | val_mask[:, i] | test_mask[:, i]
          ).sum() == num_nodes
  
      transform = RandomNodeSplit(
          split="random",
          num_splits=num_splits,
          num_train_per_class=10,
          num_val=100,
          num_test=200,
      )
      assert str(transform) == "RandomNodeSplit(split=random)"
      data = transform(data)
      assert len(data) == 5
  
      train_mask = data.train_mask
      train_mask = (
          train_mask.unsqueeze(-1) if num_splits == 1 else train_mask
      )
      val_mask = data.val_mask
      val_mask = val_mask.unsqueeze(-1) if num_splits == 1 else val_mask
      test_mask = data.test_mask
      test_mask = (
          test_mask.unsqueeze(-1) if num_splits == 1 else test_mask
      )
  
      for i in range(train_mask.size(-1)):
          assert train_mask[:, i].sum() == 10 * num_classes
          assert val_mask[:, i].sum() == 100
          assert test_mask[:, i].sum() == 200
          assert (
              train_mask[:, i] & val_mask[:, i] & test_mask[:, i]
          ).sum() == 0
          assert (
              train_mask[:, i] | val_mask[:, i] | test_mask[:, i]
          ).sum() == 10 * num_classes + 100 + 200
  
      transform = RandomNodeSplit(
          split="random",
          num_splits=num_splits,
          num_train_per_class=10,
          num_val=0.1,
          num_test=0.2,
      )
      assert str(transform) == "RandomNodeSplit(split=random)"
      data = transform(data)
  
      train_mask = data.train_mask
      train_mask = (
          train_mask.unsqueeze(-1) if num_splits == 1 else train_mask
      )
      val_mask = data.val_mask
      val_mask = val_mask.unsqueeze(-1) if num_splits == 1 else val_mask
      test_mask = data.test_mask
      test_mask = (
          test_mask.unsqueeze(-1) if num_splits == 1 else test_mask
      )
  
      for i in range(train_mask.size(-1)):
          assert train_mask[:, i].sum() == 10 * num_classes
          assert val_mask[:, i].sum() == 100
          assert test_mask[:, i].sum() == 200
          assert (
              train_mask[:, i] & val_mask[:, i] & test_mask[:, i]
          ).sum() == 0
          assert (
              train_mask[:, i] | val_mask[:, i] | test_mask[:, i]
          ).sum() == 10 * num_classes + 100 + 200
  
imports: import pytest
  import torch
  from torch_geometric.data import Data, HeteroData
  from torch_geometric.transforms import RandomNodeSplit
label_desc: Use image node_id 1 for calling a global function with example usage: test_random_node_split(num_splits) without return types