Unverified Commit a38bb5d1 authored by Hongzhi (Steve), Chen's avatar Hongzhi (Steve), Chen Committed by GitHub
Browse files

[Misc] Better naming. (#6206)


Co-authored-by: default avatarUbuntu <ubuntu@ip-172-31-28-63.ap-northeast-1.compute.internal>
parent a438565e
...@@ -27,37 +27,37 @@ def test_torch_based_feature(in_memory): ...@@ -27,37 +27,37 @@ def test_torch_based_feature(in_memory):
a = to_on_disk_tensor(test_dir, "a", a) a = to_on_disk_tensor(test_dir, "a", a)
b = to_on_disk_tensor(test_dir, "b", b) b = to_on_disk_tensor(test_dir, "b", b)
feat_store_a = gb.TorchBasedFeature(a) feature_a = gb.TorchBasedFeature(a)
feat_store_b = gb.TorchBasedFeature(b) feature_b = gb.TorchBasedFeature(b)
assert torch.equal(feat_store_a.read(), torch.tensor([1, 2, 3])) assert torch.equal(feature_a.read(), torch.tensor([1, 2, 3]))
assert torch.equal( assert torch.equal(
feat_store_b.read(), torch.tensor([[1, 2, 3], [4, 5, 6]]) feature_b.read(), torch.tensor([[1, 2, 3], [4, 5, 6]])
) )
assert torch.equal( assert torch.equal(
feat_store_a.read(torch.tensor([0, 2])), feature_a.read(torch.tensor([0, 2])),
torch.tensor([1, 3]), torch.tensor([1, 3]),
) )
assert torch.equal( assert torch.equal(
feat_store_a.read(torch.tensor([1, 1])), feature_a.read(torch.tensor([1, 1])),
torch.tensor([2, 2]), torch.tensor([2, 2]),
) )
assert torch.equal( assert torch.equal(
feat_store_b.read(torch.tensor([1])), feature_b.read(torch.tensor([1])),
torch.tensor([[4, 5, 6]]), torch.tensor([[4, 5, 6]]),
) )
feat_store_a.update(torch.tensor([0, 1, 2]), torch.tensor([0, 1, 2])) feature_a.update(torch.tensor([0, 1, 2]), torch.tensor([0, 1, 2]))
assert torch.equal(feat_store_a.read(), torch.tensor([0, 1, 2])) assert torch.equal(feature_a.read(), torch.tensor([0, 1, 2]))
feat_store_a.update(torch.tensor([2, 0]), torch.tensor([0, 2])) feature_a.update(torch.tensor([2, 0]), torch.tensor([0, 2]))
assert torch.equal(feat_store_a.read(), torch.tensor([2, 1, 0])) assert torch.equal(feature_a.read(), torch.tensor([2, 1, 0]))
with pytest.raises(IndexError): with pytest.raises(IndexError):
feat_store_a.read(torch.tensor([0, 1, 2, 3])) feature_a.read(torch.tensor([0, 1, 2, 3]))
# For windows, the file is locked by the numpy.load. We need to delete # For windows, the file is locked by the numpy.load. We need to delete
# it before closing the temporary directory. # it before closing the temporary directory.
a = b = None a = b = None
feat_store_a = feat_store_b = None feature_a = feature_b = None
def write_tensor_to_disk(dir, name, t, fmt="torch"): def write_tensor_to_disk(dir, name, t, fmt="torch"):
...@@ -77,7 +77,7 @@ def test_torch_based_feature_store(in_memory): ...@@ -77,7 +77,7 @@ def test_torch_based_feature_store(in_memory):
b = torch.tensor([2, 5, 3]) b = torch.tensor([2, 5, 3])
write_tensor_to_disk(test_dir, "a", a, fmt="torch") write_tensor_to_disk(test_dir, "a", a, fmt="torch")
write_tensor_to_disk(test_dir, "b", b, fmt="numpy") write_tensor_to_disk(test_dir, "b", b, fmt="numpy")
feat_data = [ feature_data = [
gb.OnDiskFeatureData( gb.OnDiskFeatureData(
domain="node", domain="node",
type="paper", type="paper",
...@@ -95,19 +95,19 @@ def test_torch_based_feature_store(in_memory): ...@@ -95,19 +95,19 @@ def test_torch_based_feature_store(in_memory):
in_memory=in_memory, in_memory=in_memory,
), ),
] ]
feat_store = gb.TorchBasedFeatureStore(feat_data) feature_store = gb.TorchBasedFeatureStore(feature_data)
assert torch.equal( assert torch.equal(
feat_store.read("node", "paper", "a"), torch.tensor([1, 2, 3]) feature_store.read("node", "paper", "a"), torch.tensor([1, 2, 3])
) )
assert torch.equal( assert torch.equal(
feat_store.read("edge", "paper-cites-paper", "b"), feature_store.read("edge", "paper-cites-paper", "b"),
torch.tensor([2, 5, 3]), torch.tensor([2, 5, 3]),
) )
# For windows, the file is locked by the numpy.load. We need to delete # For windows, the file is locked by the numpy.load. We need to delete
# it before closing the temporary directory. # it before closing the temporary directory.
a = b = None a = b = None
feat_stores = None feature_stores = None
# ``domain`` should be enum. # ``domain`` should be enum.
with pytest.raises(pydantic.ValidationError): with pytest.raises(pydantic.ValidationError):
...@@ -121,7 +121,7 @@ def test_torch_based_feature_store(in_memory): ...@@ -121,7 +121,7 @@ def test_torch_based_feature_store(in_memory):
) )
# ``type`` could be null. # ``type`` could be null.
feat_data = [ feature_data = [
gb.OnDiskFeatureData( gb.OnDiskFeatureData(
domain="node", domain="node",
name="a", name="a",
...@@ -130,8 +130,8 @@ def test_torch_based_feature_store(in_memory): ...@@ -130,8 +130,8 @@ def test_torch_based_feature_store(in_memory):
in_memory=True, in_memory=True,
), ),
] ]
feat_store = gb.TorchBasedFeatureStore(feat_data) feature_store = gb.TorchBasedFeatureStore(feature_data)
assert torch.equal( assert torch.equal(
feat_store.read("node", None, "a"), torch.tensor([1, 2, 3]) feature_store.read("node", None, "a"), torch.tensor([1, 2, 3])
) )
feat_stores = None feature_stores = None
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment