Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Test sparse GAT properly #1781

Draft
wants to merge 1 commit into
base: bugfix/1251-test-full-batch-saving
Choose a base branch
from
Draft

Test sparse GAT properly #1781

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
19 changes: 12 additions & 7 deletions tests/layer/test_graph_attention.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
from tensorflow import keras
from tensorflow.keras import backend as K
from tensorflow.keras.layers import Input
from stellargraph import custom_keras_layers
from stellargraph.mapper import (
FullBatchNodeGenerator,
FullBatchLinkGenerator,
Expand Down Expand Up @@ -424,7 +425,11 @@ def test_gat_build_constructor(self):
x_in, x_out = gat.in_out_tensors()
assert len(x_in) == 4 if self.sparse else 3
assert int(x_in[0].shape[-1]) == self.F_in
assert K.int_shape(x_in[-1]) == (1, G.number_of_nodes(), G.number_of_nodes())
if self.sparse:
assert K.int_shape(x_in[-2]) == (1, None, 2)
assert K.int_shape(x_in[-1]) == (1, None)
else:
assert K.int_shape(x_in[-1]) == (1, G.number_of_nodes(), G.number_of_nodes())
assert int(x_out.shape[-1]) == self.layer_sizes[-1]

def test_gat_build_linkmodel_constructor(self):
Expand All @@ -445,6 +450,9 @@ def test_gat_build_linkmodel_constructor(self):
assert int(x_out.shape[-1]) == self.layer_sizes[-1]

def test_gat_build_constructor_no_generator(self):
if self.sparse:
pytest.skip("no_generator is always dense")

G = example_graph(feature_size=self.F_in)
gat = GAT(
layer_sizes=self.layer_sizes,
Expand All @@ -458,7 +466,7 @@ def test_gat_build_constructor_no_generator(self):
assert gat.use_sparse == False

x_in, x_out = gat.in_out_tensors()
assert len(x_in) == 4 if self.sparse else 3
assert len(x_in) == 3
assert int(x_in[0].shape[-1]) == self.F_in
assert int(x_out.shape[-1]) == self.layer_sizes[-1]

Expand Down Expand Up @@ -570,10 +578,7 @@ def test_gat_serialize(self):
# Load model from json & set all weights
model2 = keras.models.model_from_json(
model_json,
custom_objects={
"GraphAttention": GraphAttention,
"GatherIndices": GatherIndices,
},
custom_objects=custom_keras_layers,
)
model2.set_weights(model_weights)

Expand Down Expand Up @@ -622,6 +627,6 @@ def test_save_load(self, tmpdir):
test_utils.model_save_load(tmpdir, gat)


def TestGATsparse(Test_GAT):
class TestGATsparse(Test_GAT):
sparse = True
method = "gat"