code: def test_lightning_hetero_link_data_custom_store():
      torch.manual_seed(12345)
  
      feature_store = MyFeatureStore()
      graph_store = MyGraphStore()
  
      x = torch.arange(10)
      feature_store.put_tensor(
          x, group_name="paper", attr_name="x", index=None
      )
      feature_store.put_tensor(
          x, group_name="author", attr_name="x", index=None
      )
      feature_store.put_tensor(
          x, group_name="term", attr_name="x", index=None
      )
  
      edge_index = get_random_edge_index(10, 10, 10)
      graph_store.put_edge_index(
          edge_index=(edge_index[0], edge_index[1]),
          edge_type=("paper", "to", "author"),
          layout="coo",
          size=(10, 10),
      )
      graph_store.put_edge_index(
          edge_index=(edge_index[0], edge_index[1]),
          edge_type=("author", "to", "paper"),
          layout="coo",
          size=(10, 10),
      )
      graph_store.put_edge_index(
          edge_index=(edge_index[0], edge_index[1]),
          edge_type=("paper", "to", "term"),
          layout="coo",
          size=(10, 10),
      )
  
      datamodule = LightningLinkData(
          (feature_store, graph_store),
          input_train_edges=("author", "to", "paper"),
          loader="neighbor",
          num_neighbors=[5],
          batch_size=32,
          num_workers=0,
      )
  
      batch = next(iter(datamodule.train_dataloader()))
      assert "edge_label_index" in batch["author", "paper"]
  
imports: import math
  import pytest
  import torch
  import torch.nn.functional
  from torch import Tensor
  from torch_geometric.data import Data, HeteroData
  from torch_geometric.data.lightning import LightningDataset, LightningLinkData, LightningNodeData
  from torch_geometric.nn import global_mean_pool
  from torch_geometric.sampler import BaseSampler, NeighborSampler
  from torch_geometric.testing import MyFeatureStore, MyGraphStore, get_random_edge_index, onlyCUDA, onlyFullTest, onlyNeighborSampler, onlyOnline, withPackage
label_desc: Use image node_id 6 for calling a global function with example usage: test_lightning_hetero_link_data_custom_store() without return types