code: def test_hetero_batch():
      e1 = ("p", "a")
      e2 = ("a", "p")
      data1 = HeteroData()
      data1["p"].x = torch.randn(100, 128)
      data1["a"].x = torch.randn(200, 128)
      data1[e1].edge_index = get_random_edge_index(100, 200, 500)
      data1[e1].edge_attr = torch.randn(500, 32)
      data1[e2].edge_index = get_random_edge_index(200, 100, 400)
      data1[e2].edge_attr = torch.randn(400, 32)
  
      data2 = HeteroData()
      data2["p"].x = torch.randn(50, 128)
      data2["a"].x = torch.randn(100, 128)
      data2[e1].edge_index = get_random_edge_index(50, 100, 300)
      data2[e1].edge_attr = torch.randn(300, 32)
      data2[e2].edge_index = get_random_edge_index(100, 50, 200)
      data2[e2].edge_attr = torch.randn(200, 32)
  
      batch = Batch.from_data_list([data1, data2])
  
      assert batch.num_graphs == len(batch) == 2
      assert batch.num_nodes == 450
  
      assert torch.allclose(batch["p"].x[:100], data1["p"].x)
      assert torch.allclose(batch["a"].x[:200], data1["a"].x)
      assert torch.allclose(batch["p"].x[100:], data2["p"].x)
      assert torch.allclose(batch["a"].x[200:], data2["a"].x)
      assert (
          batch[e1].edge_index.tolist()
          == torch.cat(
              [
                  data1[e1].edge_index,
                  data2[e1].edge_index + torch.tensor([[100], [200]]),
              ],
              1,
          ).tolist()
      )
      assert torch.allclose(
          batch[e1].edge_attr,
          torch.cat([data1[e1].edge_attr, data2[e1].edge_attr], 0),
      )
      assert (
          batch[e2].edge_index.tolist()
          == torch.cat(
              [
                  data1[e2].edge_index,
                  data2[e2].edge_index + torch.tensor([[200], [100]]),
              ],
              1,
          ).tolist()
      )
      assert torch.allclose(
          batch[e2].edge_attr,
          torch.cat([data1[e2].edge_attr, data2[e2].edge_attr], 0),
      )
      assert batch["p"].batch.size() == (150,)
      assert batch["p"].ptr.size() == (3,)
      assert batch["a"].batch.size() == (300,)
      assert batch["a"].ptr.size() == (3,)
  
      out1 = batch[0]
      assert len(out1) == 3
      assert out1.num_nodes == 300
      assert torch.allclose(out1["p"].x, data1["p"].x)
      assert torch.allclose(out1["a"].x, data1["a"].x)
      assert (
          out1[e1].edge_index.tolist() == data1[e1].edge_index.tolist()
      )
      assert torch.allclose(out1[e1].edge_attr, data1[e1].edge_attr)
      assert (
          out1[e2].edge_index.tolist() == data1[e2].edge_index.tolist()
      )
      assert torch.allclose(out1[e2].edge_attr, data1[e2].edge_attr)
  
      out2 = batch[1]
      assert len(out2) == 3
      assert out2.num_nodes == 150
      assert torch.allclose(out2["p"].x, data2["p"].x)
      assert torch.allclose(out2["a"].x, data2["a"].x)
      assert (
          out2[e1].edge_index.tolist() == data2[e1].edge_index.tolist()
      )
      assert torch.allclose(out2[e1].edge_attr, data2[e1].edge_attr)
      assert (
          out2[e2].edge_index.tolist() == data2[e2].edge_index.tolist()
      )
      assert torch.allclose(out2[e2].edge_attr, data2[e2].edge_attr)
  
imports: import os.path
  import numpy
  import pytest
  import torch
  import torch_geometric
  from torch_geometric import EdgeIndex
  from torch_geometric.data import Batch, Data, HeteroData
  from torch_geometric.testing import get_random_edge_index, withPackage
  from torch_geometric.typing import SparseTensor
  from torch_geometric.utils import to_edge_index, to_torch_sparse_tensor
label_desc: Use image node_id 9 for calling a global function with example usage: test_hetero_batch() without return types