def test_attn_agg_apply(): agg = AttentionalAggregator(2, bias=False, act="linear", kernel_initializer="ones") agg.attn_act = keras.activations.get("linear") # Self features inp1 = keras.Input(shape=(1, 2)) # Neighbour features inp2 = keras.Input(shape=(1, 2, 2)) out = agg([inp1, inp2]) # The AttentionalAggregator implmentation is a hack at the moment, it doesn't # assign any dimensions in the output to head-node features. assert agg.weight_dims == [0, 2] # Numerical test values x1 = np.array([[[1, 1]]]) x2 = np.array([[[[2, 2], [3, 3]]]]) # Agg output: # hs = relu(x1 · ones(2x2)) = [2,2] # hn = relu(x2 · ones(2x2)) = [[2,2], [4,4], [6,6]] # attn_u = ones(2) · hs + ones(2) · hn = [8, 12, 16] # attn = softmax(attn_u) = [3.3e-4, 1.8e-4, 9.81e-1] # hout = attn · hn = [5.96, 5.96] model = keras.Model(inputs=[inp1, inp2], outputs=out) actual = model.predict([x1, x2]) expected = np.array([[[5.963, 5.963]]]) assert expected == pytest.approx(actual, rel=1e-4)
def test_attn_agg_constructor(): agg = AttentionalAggregator(2, bias=False) assert agg.output_dim == 2 assert not agg.has_bias assert agg.act.__name__ == "relu" # assert agg.attn_act.__name__ == "relu" # Check config config = agg.get_config() assert config["output_dim"] == 2 assert config["bias"] is False assert config["act"] == "relu"
def test_attn_agg_zero_neighbours(): agg = AttentionalAggregator(4, bias=False, act="linear", kernel_initializer="ones") inp1 = keras.Input(shape=(1, 2)) inp2 = keras.Input(shape=(1, 0, 2)) out = agg([inp1, inp2]) model = keras.Model(inputs=[inp1, inp2], outputs=out) x1 = np.array([[[1, 1]]]) x2 = np.zeros((1, 1, 0, 2)) actual = model.predict([x1, x2]) expected = np.array([[[2, 2, 2, 2]]]) assert expected == pytest.approx(actual)
def test_attn_agg_constructor_1(): agg = AttentionalAggregator(output_dim=4, bias=True, act=lambda x: x + 1) assert agg.output_dim == 4 assert agg.has_bias assert agg.act(2) == 3