k3-node 1.0.0__py3-none-any.whl

k3_node/layers/dense/test_dense.py

@@ -0,0 +1,377 @@
+import math
+import numpy as np
+import pytest
+from keras import layers, ops, random, Sequential
+
+from k3_node.layers.dense import (
+    Linear,
+    HeteroLinear,
+    HeteroDictLinear,
+    DenseGCNConv,
+    DenseGINConv,
+    DenseGraphConv,
+    DenseSAGEConv,
+    DenseGATConv,
+    dense_diff_pool,
+    dense_mincut_pool,
+    DMoNPooling,
+)
+
+
+@pytest.mark.parametrize('weight', ['glorot', 'kaiming_uniform', None])
+@pytest.mark.parametrize('bias', ['zeros', None])
+def test_linear(weight, bias):
+    x = random.normal((3, 4, 16))
+    lin = Linear(16, 32, weight_initializer=weight, bias_initializer=bias)
+    assert str(lin) == 'Linear(16, 32, bias=True)'
+    out = lin(x)
+    assert ops.shape(out) == (3, 4, 32)
+
+
+@pytest.mark.parametrize('weight', ['glorot', 'kaiming_uniform', None])
+@pytest.mark.parametrize('bias', ['zeros', None])
+def test_lazy_linear(weight, bias):
+    x = random.normal((3, 4, 16))
+    lin = Linear(-1, 32, weight_initializer=weight, bias_initializer=bias)
+    assert str(lin) == 'Linear(-1, 32, bias=True)'
+    out = lin(x)
+    assert ops.shape(out) == (3, 4, 32)
+    assert str(lin) == 'Linear(16, 32, bias=True)'
+
+
+def test_hetero_linear_basic():
+    x = random.normal((3, 16))
+    type_vec = ops.convert_to_tensor(np.array([0, 1, 2], dtype=np.int32))
+
+    lin = HeteroLinear(16, 32, num_types=3)
+    assert str(lin) == 'HeteroLinear(16, 32, num_types=3, bias=True)'
+
+    out = lin(x, type_vec)
+    assert ops.shape(out) == (3, 32)
+
+
+def test_lazy_hetero_linear():
+    x = random.normal((3, 16))
+    type_vec = ops.convert_to_tensor(np.array([0, 1, 2], dtype=np.int32))
+
+    lin = HeteroLinear(-1, 32, num_types=3)
+    assert str(lin) == 'HeteroLinear(-1, 32, num_types=3, bias=True)'
+
+    out = lin(x, type_vec)
+    assert ops.shape(out) == (3, 32)
+
+
+@pytest.mark.parametrize('bias', [True, False])
+def test_hetero_dict_linear(bias):
+    x_dict = {
+        'v': random.normal((3, 16)),
+        'w': random.normal((2, 8)),
+    }
+
+    lin = HeteroDictLinear({'v': 16, 'w': 8}, 32, bias=bias)
+    assert str(lin) == f"HeteroDictLinear({{'v': 16, 'w': 8}}, 32, bias={bias})"
+
+    out_dict = lin(x_dict)
+    assert len(out_dict) == 2
+    assert ops.shape(out_dict['v']) == (3, 32)
+    assert ops.shape(out_dict['w']) == (2, 32)
+
+    x_dict2 = {
+        'v': random.normal((3, 16)),
+        'w': random.normal((2, 16)),
+    }
+
+    lin2 = HeteroDictLinear(16, 32, types=['v', 'w'], bias=bias)
+    assert str(lin2) == f"HeteroDictLinear({{'v': 16, 'w': 16}}, 32, bias={bias})"
+
+    out_dict2 = lin2(x_dict2)
+    assert len(out_dict2) == 2
+    assert ops.shape(out_dict2['v']) == (3, 32)
+    assert ops.shape(out_dict2['w']) == (2, 32)
+
+
+def test_lazy_hetero_dict_linear():
+    x_dict = {
+        'v': random.normal((3, 16)),
+        'w': random.normal((2, 8)),
+    }
+
+    lin = HeteroDictLinear(-1, 32, types=['v', 'w'])
+    assert str(lin) == "HeteroDictLinear({'v': -1, 'w': -1}, 32, bias=True)"
+
+    out_dict = lin(x_dict)
+    assert len(out_dict) == 2
+    assert ops.shape(out_dict['v']) == (3, 32)
+    assert ops.shape(out_dict['w']) == (2, 32)
+
+
+def test_dense_gcn_conv():
+    channels = 16
+    conv = DenseGCNConv(channels, channels)
+    assert str(conv) == 'DenseGCNConv(16, 16)'
+
+    x = random.normal((2, 3, channels))
+    adj = ops.convert_to_tensor(np.array([
+        [
+            [0.0, 1.0, 1.0],
+            [1.0, 0.0, 1.0],
+            [1.0, 1.0, 0.0],
+        ],
+        [
+            [0.0, 1.0, 0.0],
+            [1.0, 0.0, 0.0],
+            [0.0, 0.0, 0.0],
+        ],
+    ], dtype=np.float32))
+    mask = ops.convert_to_tensor(np.array([[1, 1, 1], [1, 1, 0]], dtype=bool))
+
+    out = conv(x, adj, mask)
+    assert ops.shape(out) == (2, 3, channels)
+    assert float(ops.convert_to_numpy(ops.sum(ops.abs(out[1, 2])))) == 0.0
+
+
+def test_dense_gcn_conv_with_broadcasting():
+    batch_size, num_nodes, channels = 8, 3, 16
+    conv = DenseGCNConv(channels, channels)
+
+    x = random.normal((batch_size, num_nodes, channels))
+    adj = ops.convert_to_tensor(np.array([
+        [0.0, 1.0, 1.0],
+        [1.0, 0.0, 1.0],
+        [1.0, 1.0, 0.0],
+    ], dtype=np.float32))
+
+    assert ops.shape(conv(x, adj)) == (batch_size, num_nodes, channels)
+    mask = ops.convert_to_tensor(np.array([1, 1, 1], dtype=bool))
+    assert ops.shape(conv(x, adj, mask)) == (batch_size, num_nodes, channels)
+
+
+def test_dense_gin_conv():
+    channels = 16
+    nn = Sequential([
+        layers.Dense(channels, activation='relu'),
+        layers.Dense(channels),
+    ])
+    dense_conv = DenseGINConv(nn)
+    assert str(dense_conv) == f'DenseGINConv(nn={nn})'
+
+    x = random.normal((2, 3, channels))
+    adj = ops.convert_to_tensor(np.array([
+        [
+            [0.0, 1.0, 1.0],
+            [1.0, 0.0, 1.0],
+            [1.0, 1.0, 0.0],
+        ],
+        [
+            [0.0, 1.0, 0.0],
+            [1.0, 0.0, 0.0],
+            [0.0, 0.0, 0.0],
+        ],
+    ], dtype=np.float32))
+    mask = ops.convert_to_tensor(np.array([[1, 1, 1], [1, 1, 0]], dtype=bool))
+
+    out = dense_conv(x, adj, mask)
+    assert ops.shape(out) == (2, 3, channels)
+    assert float(ops.convert_to_numpy(ops.sum(ops.abs(out[1, 2])))) == 0.0
+
+
+def test_dense_gin_conv_with_broadcasting():
+    batch_size, num_nodes, channels = 8, 3, 16
+    nn = Sequential([
+        layers.Dense(channels, activation='relu'),
+        layers.Dense(channels),
+    ])
+    conv = DenseGINConv(nn)
+
+    x = random.normal((batch_size, num_nodes, channels))
+    adj = ops.convert_to_tensor(np.array([
+        [0.0, 1.0, 1.0],
+        [1.0, 0.0, 1.0],
+        [1.0, 1.0, 0.0],
+    ], dtype=np.float32))
+
+    assert ops.shape(conv(x, adj)) == (batch_size, num_nodes, channels)
+    mask = ops.convert_to_tensor(np.array([1, 1, 1], dtype=bool))
+    assert ops.shape(conv(x, adj, mask)) == (batch_size, num_nodes, channels)
+
+
+@pytest.mark.parametrize('aggr', ['add', 'mean', 'max'])
+def test_dense_graph_conv(aggr):
+    channels = 16
+    conv = DenseGraphConv(channels, channels, aggr=aggr)
+    assert str(conv) == 'DenseGraphConv(16, 16)'
+
+    x = random.normal((2, 3, channels))
+    adj = ops.convert_to_tensor(np.array([
+        [
+            [0.0, 1.0, 1.0],
+            [1.0, 0.0, 1.0],
+            [1.0, 1.0, 0.0],
+        ],
+        [
+            [0.0, 1.0, 0.0],
+            [1.0, 0.0, 0.0],
+            [0.0, 0.0, 0.0],
+        ],
+    ], dtype=np.float32))
+    mask = ops.convert_to_tensor(np.array([[1, 1, 1], [1, 1, 0]], dtype=bool))
+
+    out = conv(x, adj, mask)
+    assert ops.shape(out) == (2, 3, channels)
+    assert float(ops.convert_to_numpy(ops.sum(ops.abs(out[1, 2])))) == 0.0
+
+
+@pytest.mark.parametrize('aggr', ['add', 'mean', 'max'])
+def test_dense_graph_conv_with_broadcasting(aggr):
+    batch_size, num_nodes, channels = 8, 3, 16
+    conv = DenseGraphConv(channels, channels, aggr=aggr)
+
+    x = random.normal((batch_size, num_nodes, channels))
+    adj = ops.convert_to_tensor(np.array([
+        [0.0, 1.0, 1.0],
+        [1.0, 0.0, 1.0],
+        [1.0, 1.0, 0.0],
+    ], dtype=np.float32))
+
+    assert ops.shape(conv(x, adj)) == (batch_size, num_nodes, channels)
+    mask = ops.convert_to_tensor(np.array([1, 1, 1], dtype=bool))
+    assert ops.shape(conv(x, adj, mask)) == (batch_size, num_nodes, channels)
+
+
+@pytest.mark.parametrize('normalize', [True, False])
+def test_dense_sage_conv(normalize):
+    channels = 16
+    conv = DenseSAGEConv(channels, channels, normalize=normalize)
+    assert str(conv) == 'DenseSAGEConv(16, 16)'
+
+    x = random.normal((2, 3, channels))
+    adj = ops.convert_to_tensor(np.array([
+        [
+            [0.0, 1.0, 1.0],
+            [1.0, 0.0, 1.0],
+            [1.0, 1.0, 0.0],
+        ],
+        [
+            [0.0, 1.0, 0.0],
+            [1.0, 0.0, 0.0],
+            [0.0, 0.0, 0.0],
+        ],
+    ], dtype=np.float32))
+    mask = ops.convert_to_tensor(np.array([[1, 1, 1], [1, 1, 0]], dtype=bool))
+
+    out = conv(x, adj, mask)
+    assert ops.shape(out) == (2, 3, channels)
+    assert float(ops.convert_to_numpy(ops.sum(ops.abs(out[1, 2])))) == 0.0
+
+
+def test_dense_sage_conv_with_broadcasting():
+    batch_size, num_nodes, channels = 8, 3, 16
+    conv = DenseSAGEConv(channels, channels)
+
+    x = random.normal((batch_size, num_nodes, channels))
+    adj = ops.convert_to_tensor(np.array([
+        [0.0, 1.0, 1.0],
+        [1.0, 0.0, 1.0],
+        [1.0, 1.0, 0.0],
+    ], dtype=np.float32))
+
+    assert ops.shape(conv(x, adj)) == (batch_size, num_nodes, channels)
+    mask = ops.convert_to_tensor(np.array([1, 1, 1], dtype=bool))
+    assert ops.shape(conv(x, adj, mask)) == (batch_size, num_nodes, channels)
+
+
+@pytest.mark.parametrize('heads', [1, 4])
+@pytest.mark.parametrize('concat', [True, False])
+def test_dense_gat_conv(heads, concat):
+    channels = 16
+    conv = DenseGATConv(channels, channels, heads=heads, concat=concat)
+    assert str(conv) == f'DenseGATConv(16, 16, heads={heads})'
+
+    x = random.normal((2, 3, channels))
+    adj = ops.convert_to_tensor(np.array([
+        [
+            [0.0, 1.0, 0.0],
+            [1.0, 0.0, 1.0],
+            [0.0, 1.0, 0.0],
+        ],
+        [
+            [0.0, 1.0, 0.0],
+            [1.0, 0.0, 0.0],
+            [0.0, 0.0, 0.0],
+        ],
+    ], dtype=np.float32))
+    mask = ops.convert_to_tensor(np.array([[1, 1, 1], [1, 1, 0]], dtype=bool))
+
+    out = conv(x, adj, mask)
+    out_dim = heads * channels if concat else channels
+    assert ops.shape(out) == (2, 3, out_dim)
+    assert float(ops.convert_to_numpy(ops.sum(ops.abs(out[1, 2])))) == 0.0
+
+
+def test_dense_gat_conv_with_broadcasting():
+    batch_size, num_nodes, channels = 8, 3, 16
+    conv = DenseGATConv(channels, channels, heads=4)
+
+    x = random.normal((batch_size, num_nodes, channels))
+    adj = ops.convert_to_tensor(np.array([
+        [0.0, 1.0, 1.0],
+        [1.0, 0.0, 1.0],
+        [1.0, 1.0, 0.0],
+    ], dtype=np.float32))
+
+    assert ops.shape(conv(x, adj)) == (batch_size, num_nodes, 64)
+    mask = ops.convert_to_tensor(np.array([1, 1, 1], dtype=bool))
+    assert ops.shape(conv(x, adj, mask)) == (batch_size, num_nodes, 64)
+
+
+def test_dense_diff_pool():
+    batch_size, num_nodes, channels, num_clusters = (2, 20, 16, 10)
+    x = random.normal((batch_size, num_nodes, channels))
+    adj = ops.convert_to_tensor(np.random.rand(batch_size, num_nodes, num_nodes).astype(np.float32))
+    s = random.normal((batch_size, num_nodes, num_clusters))
+    mask = ops.convert_to_tensor(np.random.randint(0, 2, (batch_size, num_nodes), dtype=bool))
+
+    x_out, adj_out, link_loss, ent_loss = dense_diff_pool(x, adj, s, mask)
+    assert ops.shape(x_out) == (2, 10, 16)
+    assert ops.shape(adj_out) == (2, 10, 10)
+    assert float(ops.convert_to_numpy(link_loss)) >= 0
+    assert float(ops.convert_to_numpy(ent_loss)) >= 0
+
+
+def test_dense_mincut_pool():
+    batch_size, num_nodes, channels, num_clusters = (2, 20, 16, 10)
+    x = random.normal((batch_size, num_nodes, channels))
+    adj = ops.ones((batch_size, num_nodes, num_nodes), dtype="float32")
+    s = random.normal((batch_size, num_nodes, num_clusters))
+    mask = ops.convert_to_tensor(np.random.randint(0, 2, (batch_size, num_nodes), dtype=bool))
+
+    x_out, adj_out, mincut_loss, ortho_loss = dense_mincut_pool(x, adj, s, mask)
+    assert ops.shape(x_out) == (2, 10, 16)
+    assert ops.shape(adj_out) == (2, 10, 10)
+    assert -1.0 <= float(ops.convert_to_numpy(mincut_loss)) <= 0.0
+    assert 0.0 <= float(ops.convert_to_numpy(ortho_loss)) <= 2.0
+
+
+def test_dmon_pooling():
+    batch_size, num_nodes, channels, num_clusters = (2, 20, 16, 10)
+    x = random.normal((batch_size, num_nodes, channels))
+    adj = ops.ones((batch_size, num_nodes, num_nodes), dtype="float32")
+    mask = ops.convert_to_tensor(np.random.randint(0, 2, (batch_size, num_nodes), dtype=bool))
+
+    pool = DMoNPooling([channels, channels], num_clusters)
+    assert str(pool) == 'DMoNPooling(16, num_clusters=10)'
+
+    s, x_out, adj_out, spectral_loss, ortho_loss, cluster_loss = pool(x, adj, mask)
+    assert ops.shape(s) == (2, 20, 10)
+    assert ops.shape(x_out) == (2, 10, 16)
+    assert ops.shape(adj_out) == (2, 10, 10)
+    assert -1.0 <= float(ops.convert_to_numpy(spectral_loss)) <= 0.5
+    assert 0.0 <= float(ops.convert_to_numpy(ortho_loss)) <= math.sqrt(2) + 1e-4
+    assert 0.0 <= float(ops.convert_to_numpy(cluster_loss)) <= math.sqrt(num_clusters) - 1 + 1e-4
+
+
+
+def test_dense_gat_conv_attention_vectors_initialized_independently():
+    conv = DenseGATConv(16, 8, heads=2)
+    conv(np.random.randn(1, 5, 16).astype("float32"), np.ones((1, 5, 5), dtype="float32"))
+    assert not np.allclose(ops.convert_to_numpy(conv.att_src), ops.convert_to_numpy(conv.att_dst))

k3_node/layers/functional/__init__.py

@@ -0,0 +1,13 @@
+r"""Functional operator package."""
+
+from .bro import bro
+from .edge_dropout import EdgeDropout
+from .gini import gini
+
+__all__ = [
+    "bro",
+    "EdgeDropout",
+    "gini",
+]
+
+classes = __all__

k3_node/layers/functional/bro.py

@@ -0,0 +1,49 @@
+from typing import Union
+
+from keras import ops
+
+
+def bro(x, batch, p: Union[int, float] = 2):
+    r"""The Batch Representation Orthogonality penalty from the `"Improving
+    Molecular Graph Neural Network Explainability with Orthonormalization
+    and Induced Sparsity" <https://arxiv.org/abs/2105.04854>`_ paper.
+
+    Computes a regularization for each graph representation in a mini-batch
+    according to
+
+    .. math::
+        \mathcal{L}_{\textrm{BRO}}^\mathrm{graph} =
+          || \mathbf{HH}^T - \mathbf{I}||_p
+
+    and returns an average over all graphs in the batch.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import bro
+
+        x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
+        batch = np.repeat([0, 1], 5)  # nodes 0-4 belong to graph 0, nodes 5-9 to graph 1
+
+        loss = bro(x, batch)  # Batch Representation Orthogonality penalty
+        print(tuple(loss.shape))  # (): a scalar
+        ```
+    """
+    batch_np = ops.convert_to_numpy(batch)
+    unique_ids = sorted(set(batch_np.tolist()))
+
+    total = 0.0
+    for graph_id in unique_ids:
+        mask = ops.equal(batch, graph_id)
+        where_mask = ops.where(mask)
+        idx = where_mask[0] if isinstance(where_mask, (list, tuple)) else where_mask
+        idx = ops.reshape(idx, (-1,))
+
+        x_g = ops.take(x, idx, axis=0)
+        n = ops.shape(x_g)[0]
+        eye = ops.eye(n, dtype=x_g.dtype)
+        diff = ops.matmul(x_g, ops.transpose(x_g)) - eye
+
+        total = total + ops.power(ops.sum(ops.power(ops.abs(diff), p)), 1.0 / p)
+
+    return total / len(unique_ids)

k3_node/layers/functional/edge_dropout.py

@@ -0,0 +1,55 @@
+import keras
+from keras import ops
+
+
+class EdgeDropout(keras.layers.Layer):
+    r"""Randomly drops edges during training (DropEdge, `Rong et al., 2020
+    <https://arxiv.org/abs/1907.10903>`_), like PyG's ``dropout_edge``.
+
+    Instead of removing edges, which would change the number of edges in every step and cannot be
+    compiled (XLA / ``jax.jit``), it returns an edge weight of 0 for dropped edges and 1 for kept
+    ones. Pass it as ``edge_weight`` to layers that support edge weights (e.g. ``GCNConv``).
+
+    Args:
+        rate (float): Probability of dropping an edge. (default: ``0.5``)
+        force_undirected (bool): Drop both directions ``(i, j)`` and ``(j, i)`` of an edge together.
+            (default: :obj:`False`)
+
+    Example:
+        ```python
+        import numpy as np
+        from keras import ops
+        from k3_node.layers import EdgeDropout
+
+        edge_index = np.random.randint(0, 10, size=(2, 30))
+        drop = EdgeDropout(rate=0.2, force_undirected=True)
+        edge_weight = drop(edge_index, training=True)  # 0 for dropped edges, 1 for kept edges
+        print(tuple(edge_weight.shape))  # (30,)
+        print(float(ops.min(drop(edge_index))))  # 1.0: nothing is dropped at inference
+        ```
+    """
+
+    def __init__(self, rate: float = 0.5, force_undirected: bool = False, seed=None, **kwargs):
+        super().__init__(**kwargs)
+        self.rate = rate
+        self.force_undirected = force_undirected
+        self.seed_generator = keras.random.SeedGenerator(seed)
+
+    def call(self, edge_index, training=None):
+        edge_index = ops.cast(edge_index, "int32")
+        num_edges = ops.shape(edge_index)[1]
+        if not training or self.rate <= 0:
+            return ops.ones((num_edges,), dtype="float32")
+        if not self.force_undirected:
+            noise = keras.random.uniform((num_edges,), seed=self.seed_generator)
+        else:
+            # The same pseudo-random number for (i, j) and (j, i): hash the unordered node pair.
+            lo = ops.cast(ops.minimum(edge_index[0], edge_index[1]), "float32")
+            hi = ops.cast(ops.maximum(edge_index[0], edge_index[1]), "float32")
+            offset = keras.random.uniform((), seed=self.seed_generator) * 1000.0
+            h = ops.sin(lo * 12.9898 + hi * 78.233 + offset) * 43758.5453
+            noise = h - ops.floor(h)
+        return ops.cast(noise >= self.rate, "float32")
+
+    def get_config(self):
+        return {**super().get_config(), "rate": self.rate, "force_undirected": self.force_undirected}

k3_node/layers/functional/gini.py

@@ -0,0 +1,44 @@
+from keras import ops
+
+EPS = 1.1920929e-07  # float32 machine epsilon, matches torch.finfo().eps
+
+
+def gini(w):
+    r"""The Gini coefficient from the `"Improving Molecular Graph Neural
+    Network Explainability with Orthonormalization and Induced Sparsity"
+    <https://arxiv.org/abs/2105.04854>`_ paper.
+
+    Computes a regularization penalty :math:`\in [0, 1]` for each row of a
+    matrix according to
+
+    .. math::
+        \mathcal{L}_\textrm{Gini}^i = \sum_j^n \sum_{j'}^n \frac{|w_{ij}
+         - w_{ij'}|}{2 (n^2 - n)\bar{w_i}}
+
+    and returns an average over all rows.
+
+    Args:
+        w: A two-dimensional tensor.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import gini
+
+        w = np.random.rand(16, 8).astype("float32")  # e.g. a weight matrix
+        print(tuple(gini(w).shape))  # (): scalar sparsity measure
+        ```
+    """
+    num_rows = ops.shape(w)[0]
+    n = ops.shape(w)[-1]
+
+    total = 0.0
+    for i in range(num_rows):
+        row = ops.take(w, i, axis=0)
+        t = ops.tile(ops.reshape(row, (1, -1)), (n, 1))
+        u = ops.sum(ops.abs(t - ops.transpose(t))) / (
+            2 * (n**2 - n) * ops.mean(ops.abs(row)) + EPS
+        )
+        total = total + u
+
+    return total / num_rows

k3_node/layers/functional/test_functional.py

@@ -0,0 +1,34 @@
+from keras import ops
+
+from k3_node.layers.functional import bro, gini
+
+
+def test_bro():
+    batch = ops.convert_to_tensor([0, 0, 0, 0, 1, 1, 1, 2, 2], dtype="int64")
+
+    g1 = ops.convert_to_tensor([
+        [0.2, 0.2, 0.2, 0.2],
+        [0.0, 0.2, 0.2, 0.2],
+        [0.2, 0.0, 0.2, 0.2],
+        [0.2, 0.2, 0.0, 0.2],
+    ], dtype="float32")
+    g2 = ops.convert_to_tensor([
+        [0.2, 0.2, 0.2, 0.2],
+        [0.0, 0.2, 0.2, 0.2],
+        [0.2, 0.0, 0.2, 0.2],
+    ], dtype="float32")
+    g3 = ops.convert_to_tensor([
+        [0.2, 0.2, 0.2, 0.2],
+        [0.2, 0.0, 0.2, 0.2],
+    ], dtype="float32")
+    x = ops.concatenate([g1, g2, g3], axis=0)
+
+    out = bro(x, batch)
+    assert ops.shape(out) == ()
+    assert float(out) > 0.0
+
+
+def test_gini():
+    w = ops.convert_to_tensor([[0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 1000.0]], dtype="float32")
+    out = gini(w)
+    assert abs(float(out) - 0.5) < 1e-5

k3_node/layers/kge/__init__.py

@@ -0,0 +1,17 @@
+r"""Knowledge Graph Embedding (KGE) package."""
+
+from .base import KGEModel
+from .transe import TransE
+from .complex import ComplEx
+from .distmult import DistMult
+from .rotate import RotatE
+from .loader import KGTripletLoader
+
+__all__ = classes = [
+    "KGEModel",
+    "TransE",
+    "ComplEx",
+    "DistMult",
+    "RotatE",
+    "KGTripletLoader",
+]