code: def test_features(idtype):
      """Test the features of batched DGLGraphs"""
      g1 = dgl.heterograph(
          {
              ("user", "follows", "user"): ([0, 1], [1, 2]),
              ("user", "plays", "game"): ([0, 1], [0, 0]),
          },
          idtype=idtype,
          device=F.ctx(),
      )
      g1.nodes["user"].data["h1"] = F.tensor([[0.0], [1.0], [2.0]])
      g1.nodes["user"].data["h2"] = F.tensor([[3.0], [4.0], [5.0]])
      g1.nodes["game"].data["h1"] = F.tensor([[0.0]])
      g1.nodes["game"].data["h2"] = F.tensor([[1.0]])
      g1.edges["follows"].data["h1"] = F.tensor([[0.0], [1.0]])
      g1.edges["follows"].data["h2"] = F.tensor([[2.0], [3.0]])
      g1.edges["plays"].data["h1"] = F.tensor([[0.0], [1.0]])
  
      g2 = dgl.heterograph(
          {
              ("user", "follows", "user"): ([0, 1], [1, 2]),
              ("user", "plays", "game"): ([0, 1], [0, 0]),
          },
          idtype=idtype,
          device=F.ctx(),
      )
      g2.nodes["user"].data["h1"] = F.tensor([[0.0], [1.0], [2.0]])
      g2.nodes["user"].data["h2"] = F.tensor([[3.0], [4.0], [5.0]])
      g2.nodes["game"].data["h1"] = F.tensor([[0.0]])
      g2.nodes["game"].data["h2"] = F.tensor([[1.0]])
      g2.edges["follows"].data["h1"] = F.tensor([[0.0], [1.0]])
      g2.edges["follows"].data["h2"] = F.tensor([[2.0], [3.0]])
      g2.edges["plays"].data["h1"] = F.tensor([[0.0], [1.0]])
  
      # test default setting
      bg = dgl.batch([g1, g2])
      assert F.allclose(
          bg.nodes["user"].data["h1"],
          F.cat(
              [
                  g1.nodes["user"].data["h1"],
                  g2.nodes["user"].data["h1"],
              ],
              dim=0,
          ),
      )
      assert F.allclose(
          bg.nodes["user"].data["h2"],
          F.cat(
              [
                  g1.nodes["user"].data["h2"],
                  g2.nodes["user"].data["h2"],
              ],
              dim=0,
          ),
      )
      assert F.allclose(
          bg.nodes["game"].data["h1"],
          F.cat(
              [
                  g1.nodes["game"].data["h1"],
                  g2.nodes["game"].data["h1"],
              ],
              dim=0,
          ),
      )
      assert F.allclose(
          bg.nodes["game"].data["h2"],
          F.cat(
              [
                  g1.nodes["game"].data["h2"],
                  g2.nodes["game"].data["h2"],
              ],
              dim=0,
          ),
      )
      assert F.allclose(
          bg.edges["follows"].data["h1"],
          F.cat(
              [
                  g1.edges["follows"].data["h1"],
                  g2.edges["follows"].data["h1"],
              ],
              dim=0,
          ),
      )
      assert F.allclose(
          bg.edges["follows"].data["h2"],
          F.cat(
              [
                  g1.edges["follows"].data["h2"],
                  g2.edges["follows"].data["h2"],
              ],
              dim=0,
          ),
      )
      assert F.allclose(
          bg.edges["plays"].data["h1"],
          F.cat(
              [
                  g1.edges["plays"].data["h1"],
                  g2.edges["plays"].data["h1"],
              ],
              dim=0,
          ),
      )
  
      # test specifying ndata/edata
      bg = dgl.batch([g1, g2], ndata=["h2"], edata=["h1"])
      assert F.allclose(
          bg.nodes["user"].data["h2"],
          F.cat(
              [
                  g1.nodes["user"].data["h2"],
                  g2.nodes["user"].data["h2"],
              ],
              dim=0,
          ),
      )
      assert F.allclose(
          bg.nodes["game"].data["h2"],
          F.cat(
              [
                  g1.nodes["game"].data["h2"],
                  g2.nodes["game"].data["h2"],
              ],
              dim=0,
          ),
      )
      assert F.allclose(
          bg.edges["follows"].data["h1"],
          F.cat(
              [
                  g1.edges["follows"].data["h1"],
                  g2.edges["follows"].data["h1"],
              ],
              dim=0,
          ),
      )
      assert F.allclose(
          bg.edges["plays"].data["h1"],
          F.cat(
              [
                  g1.edges["plays"].data["h1"],
                  g2.edges["plays"].data["h1"],
              ],
              dim=0,
          ),
      )
      assert "h1" not in bg.nodes["user"].data
      assert "h1" not in bg.nodes["game"].data
      assert "h2" not in bg.edges["follows"].data
  
      # Test unbatching graphs
      g3, g4 = dgl.unbatch(bg)
      check_equivalence_between_heterographs(
          g1,
          g3,
          node_attrs={"user": ["h2"], "game": ["h2"]},
          edge_attrs={("user", "follows", "user"): ["h1"]},
      )
      check_equivalence_between_heterographs(
          g2,
          g4,
          node_attrs={"user": ["h2"], "game": ["h2"]},
          edge_attrs={("user", "follows", "user"): ["h1"]},
      )
  
file_path: dgl/tests/python/common/test_batch-heterograph.py
imports: import unittest
  import backend
  import dgl
  import pytest
  from dgl.base import ALL
  from utils import check_graph_equal, get_cases, parametrize_idtype
label_desc: Use image node_id 4 for calling a global function with example usage: test_features(idtype) without return types