code: def test_torch_based_feature_store(in_memory):
      with tempfile.TemporaryDirectory() as test_dir:
          a = torch.tensor([[1, 2, 4], [2, 5, 3]])
          b = torch.tensor([[[1, 2], [3, 4]], [[2, 5], [3, 4]]])
          write_tensor_to_disk(test_dir, "a", a, fmt="torch")
          write_tensor_to_disk(test_dir, "b", b, fmt="numpy")
          feature_data = [
              gb.OnDiskFeatureData(
                  domain="node",
                  type="paper",
                  name="a",
                  format="torch",
                  path=os.path.join(test_dir, "a.pt"),
                  in_memory=True,
              ),
              gb.OnDiskFeatureData(
                  domain="edge",
                  type="paper:cites:paper",
                  name="b",
                  format="numpy",
                  path=os.path.join(test_dir, "b.npy"),
                  in_memory=in_memory,
              ),
          ]
          feature_store = gb.TorchBasedFeatureStore(feature_data)
  
          # Test read the entire feature.
          assert torch.equal(
              feature_store.read("node", "paper", "a"),
              torch.tensor([[1, 2, 4], [2, 5, 3]]),
          )
          assert torch.equal(
              feature_store.read("edge", "paper:cites:paper", "b"),
              torch.tensor([[[1, 2], [3, 4]], [[2, 5], [3, 4]]]),
          )
  
          # Test get the size of the entire feature.
          assert feature_store.size("node", "paper", "a") == torch.Size(
              [3]
          )
          assert feature_store.size(
              "edge", "paper:cites:paper", "b"
          ) == torch.Size([2, 2])
  
          # Test get the keys of the features.
          assert feature_store.keys() == [
              ("node", "paper", "a"),
              ("edge", "paper:cites:paper", "b"),
          ]
  
          # For windows, the file is locked by the numpy.load. We need to delete
          # it before closing the temporary directory.
          a = b = None
          feature_store = None
  
          # ``domain`` should be enum.
          with pytest.raises(pydantic.ValidationError):
              _ = gb.OnDiskFeatureData(
                  domain="invalid",
                  type="paper",
                  name="a",
                  format="torch",
                  path=os.path.join(test_dir, "a.pt"),
                  in_memory=True,
              )
  
          # ``type`` could be null.
          feature_data = [
              gb.OnDiskFeatureData(
                  domain="node",
                  name="a",
                  format="torch",
                  path=os.path.join(test_dir, "a.pt"),
                  in_memory=True,
              ),
          ]
          feature_store = gb.TorchBasedFeatureStore(feature_data)
          # Test read the entire feature.
          assert torch.equal(
              feature_store.read("node", None, "a"),
              torch.tensor([[1, 2, 4], [2, 5, 3]]),
          )
          # Test get the size of the entire feature.
          assert feature_store.size("node", None, "a") == torch.Size(
              [3]
          )
  
          feature_store = None
  
imports: import os
  import tempfile
  import unittest
  from functools import reduce
  from operator import mul
  import backend
  import numpy
  import pydantic
  import pytest
  import torch
  from dgl import graphbolt
label_desc: Use image node_id 9 for calling a global function with example usage: test_torch_based_feature_store(in_memory) without return types