k3-node 1.0.0__py3-none-any.whl

k3_node/models/mole_bert.py

@@ -0,0 +1,379 @@
+import os
+import os.path as osp
+import shutil
+from typing import Optional, Tuple, Union
+
+import keras
+from keras import layers, ops
+
+from k3_node.layers.conv.message_passing import MessagePassing
+from k3_node.layers.conv.utils import add_self_loops
+from k3_node.layers.pool import global_add_pool, global_max_pool, global_mean_pool
+from k3_node.data.download import download_url
+from k3_node.ops.creation import full
+
+
+MOLE_BERT_URL = (
+    "https://github.com/junxia97/Mole-BERT/raw/refs/heads/main/model_gin/Mole-BERT.pth"
+)
+
+
+class MoleBERTGINConv(MessagePassing):
+    """Extension of GIN aggregation to incorporate categorical edge information with self-loops.
+
+    Matches the GINConv variant from Hu et al. used in Mole-BERT.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.models import MoleBERTGINConv
+
+        x = np.random.rand(4, 32).astype("float32")  # atom embeddings
+        edge_index = np.array([[0, 1, 1, 2], [1, 0, 2, 1]])
+        edge_attr = np.array([[0, 0], [0, 0], [1, 0], [1, 0]])  # bond type, bond direction
+        conv = MoleBERTGINConv(emb_dim=32)
+        print(tuple(conv(x, edge_index, edge_attr).shape))  # (4, 32)
+        ```
+    """
+
+    def __init__(
+        self,
+        emb_dim: int = 300,
+        out_dim: Optional[int] = None,
+        num_bond_type: int = 6,
+        num_bond_direction: int = 3,
+        **kwargs,
+    ):
+        super().__init__(aggr="add", **kwargs)
+        self.emb_dim = emb_dim
+        self.out_dim = emb_dim if out_dim is None else out_dim
+        self.num_bond_type = num_bond_type
+        self.num_bond_direction = num_bond_direction
+
+        self.mlp = keras.Sequential(
+            [
+                layers.Dense(2 * emb_dim, activation="relu", name="mlp_0"),
+                layers.Dense(self.out_dim, name="mlp_2"),
+            ],
+            name="mlp",
+        )
+        self.edge_embedding1 = layers.Embedding(
+            num_bond_type, emb_dim, name="edge_embedding1"
+        )
+        self.edge_embedding2 = layers.Embedding(
+            num_bond_direction, emb_dim, name="edge_embedding2"
+        )
+
+    def build(self, input_shape=None):
+        self.mlp.build((None, self.emb_dim))
+        self.edge_embedding1.build((None,))
+        self.edge_embedding2.build((None,))
+        super().build(input_shape)
+
+    def call(self, x, edge_index, edge_attr):
+        num_nodes = ops.shape(x)[0]
+
+        # Add self-loops to edge space
+        edge_index, _ = add_self_loops(edge_index, num_nodes=num_nodes)
+
+        # Add features corresponding to self-loop edges: [4, 0]
+        self_loop_attr = ops.stack(
+            [
+                full((num_nodes,), 4, dtype=edge_attr.dtype),
+                ops.zeros((num_nodes,), dtype=edge_attr.dtype),
+            ],
+            axis=1,
+        )
+        edge_attr = ops.concatenate([edge_attr, self_loop_attr], axis=0)
+
+        edge_embeddings = self.edge_embedding1(edge_attr[:, 0]) + self.edge_embedding2(
+            edge_attr[:, 1]
+        )
+
+        return self.propagate(
+            edge_index, x=x, edge_attr=edge_embeddings, size=(num_nodes, num_nodes)
+        )
+
+    def message(self, x_j, edge_attr):
+        return x_j + edge_attr
+
+    def update(self, aggr_out):
+        return self.mlp(aggr_out)
+
+
+class MoleBERTGNN(layers.Layer):
+    """5-layer GIN encoder backbone of Mole-BERT with Jumping Knowledge.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.models import MoleBERTGNN
+
+        x = np.stack([np.random.randint(0, 119, size=5), np.random.randint(0, 3, size=5)], axis=1)  # atom type, chirality
+        edge_index = np.array([[0, 1, 1, 2, 3, 4], [1, 0, 2, 1, 4, 3]])
+        edge_attr = np.stack([np.random.randint(0, 5, size=6), np.random.randint(0, 3, size=6)], axis=1)  # bond type, direction
+        batch = np.array([0, 0, 0, 1, 1])  # two molecules
+
+        gnn = MoleBERTGNN(num_layer=3, emb_dim=32, JK="last")
+        print(tuple(gnn(x, edge_index, edge_attr).shape))  # (5, 32)
+        ```
+    """
+
+    def __init__(
+        self,
+        num_layer: int = 5,
+        emb_dim: int = 300,
+        num_atom_type: int = 120,
+        num_chirality_tag: int = 3,
+        JK: str = "last",
+        drop_ratio: float = 0.0,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        if num_layer < 2:
+            raise ValueError("Number of GNN layers must be greater than 1.")
+        self.num_layer = num_layer
+        self.emb_dim = emb_dim
+        self.num_atom_type = num_atom_type
+        self.num_chirality_tag = num_chirality_tag
+        self.JK = JK
+        self.drop_ratio = drop_ratio
+
+        self.x_embedding1 = layers.Embedding(
+            num_atom_type, emb_dim, name="x_embedding1"
+        )
+        self.x_embedding2 = layers.Embedding(
+            num_chirality_tag, emb_dim, name="x_embedding2"
+        )
+
+        self.gnns = [
+            MoleBERTGINConv(emb_dim=emb_dim, name=f"gnns_{i}")
+            for i in range(num_layer)
+        ]
+        self.batch_norms = [
+            layers.BatchNormalization(
+                axis=-1, epsilon=1e-5, momentum=0.9, name=f"batch_norms_{i}"
+            )
+            for i in range(num_layer)
+        ]
+        self.dropout_layer = layers.Dropout(drop_ratio)
+
+    def build(self, input_shape=None):
+        self.x_embedding1.build((None,))
+        self.x_embedding2.build((None,))
+        for i in range(self.num_layer):
+            self.gnns[i].build(None)
+            self.batch_norms[i].build((None, self.emb_dim))
+        super().build(input_shape)
+
+    def call(self, x, edge_index, edge_attr, training: bool = False):
+        h = self.x_embedding1(x[:, 0]) + self.x_embedding2(x[:, 1])
+        h_list = [h]
+
+        for layer in range(self.num_layer):
+            h = self.gnns[layer](h_list[layer], edge_index, edge_attr)
+            h = self.batch_norms[layer](h, training=training)
+            if layer == self.num_layer - 1:
+                # Remove ReLU for the last layer
+                h = self.dropout_layer(h, training=training)
+            else:
+                h = self.dropout_layer(ops.relu(h), training=training)
+            h_list.append(h)
+
+        if self.JK == "concat":
+            node_representation = ops.concatenate(h_list, axis=1)
+        elif self.JK == "last":
+            node_representation = h_list[-1]
+        elif self.JK == "max":
+            node_representation = ops.max(ops.stack(h_list, axis=0), axis=0)
+        elif self.JK == "sum":
+            node_representation = ops.sum(ops.stack(h_list, axis=0), axis=0)
+        else:
+            raise ValueError(f"Unknown JK mode: {self.JK}")
+
+        return node_representation
+
+
+class MoleBERT(keras.Model):
+    """Complete Mole-BERT Model with graph-level pooling and property prediction head.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.models import MoleBERT
+
+        x = np.stack([np.random.randint(0, 119, size=5), np.random.randint(0, 3, size=5)], axis=1)  # atom type, chirality
+        edge_index = np.array([[0, 1, 1, 2, 3, 4], [1, 0, 2, 1, 4, 3]])
+        edge_attr = np.stack([np.random.randint(0, 5, size=6), np.random.randint(0, 3, size=6)], axis=1)  # bond type, direction
+        batch = np.array([0, 0, 0, 1, 1])  # two molecules
+
+        model = MoleBERT(num_layer=3, emb_dim=32, num_tasks=2, graph_pooling="mean")
+        logits, node_rep = model((x, edge_index, edge_attr, batch))
+        print(tuple(logits.shape), tuple(node_rep.shape))  # (2, 2) (5, 32): per-molecule predictions, per-atom embeddings
+        ```
+    """
+
+    def __init__(
+        self,
+        num_layer: int = 5,
+        emb_dim: int = 300,
+        num_tasks: Optional[int] = None,
+        JK: str = "last",
+        drop_ratio: float = 0.0,
+        graph_pooling: str = "mean",
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.num_layer = num_layer
+        self.emb_dim = emb_dim
+        self.num_tasks = num_tasks
+        self.JK = JK
+        self.drop_ratio = drop_ratio
+        self.graph_pooling = graph_pooling
+
+        self.gnn = MoleBERTGNN(
+            num_layer=num_layer,
+            emb_dim=emb_dim,
+            JK=JK,
+            drop_ratio=drop_ratio,
+            name="gnn",
+        )
+
+        if graph_pooling in ("sum", "add"):
+            self.pool_fn = global_add_pool
+        elif graph_pooling == "mean":
+            self.pool_fn = global_mean_pool
+        elif graph_pooling == "max":
+            self.pool_fn = global_max_pool
+        else:
+            raise ValueError(f"Invalid graph pooling type: '{graph_pooling}'")
+
+        if num_tasks is not None:
+            mult = (num_layer + 1) if JK == "concat" else 1
+            self.graph_pred_linear = layers.Dense(
+                num_tasks, name="graph_pred_linear"
+            )
+        else:
+            self.graph_pred_linear = None
+
+    def build(self, input_shape=None):
+        self.gnn.build(None)
+        if self.graph_pred_linear is not None:
+            mult = (self.num_layer + 1) if self.JK == "concat" else 1
+            self.graph_pred_linear.build((None, mult * self.emb_dim))
+        super().build(input_shape)
+
+    def call(self, inputs, training: bool = False):
+        """Call MoleBERT model.
+
+        inputs can be a tuple: `(x, edge_index, edge_attr)` or `(x, edge_index, edge_attr, batch)`.
+        """
+        if isinstance(inputs, (list, tuple)):
+            if len(inputs) == 3:
+                x, edge_index, edge_attr = inputs
+                batch = None
+            elif len(inputs) == 4:
+                x, edge_index, edge_attr, batch = inputs
+            else:
+                raise ValueError("Expected 3 or 4 input tensors.")
+        elif isinstance(inputs, dict):
+            x = inputs["x"]
+            edge_index = inputs["edge_index"]
+            edge_attr = inputs["edge_attr"]
+            batch = inputs.get("batch", None)
+        else:
+            raise ValueError("inputs must be a tuple, list, or dict.")
+
+        node_rep = self.gnn(x, edge_index, edge_attr, training=training)
+
+        if batch is not None:
+            graph_rep = self.pool_fn(node_rep, batch)
+            if self.graph_pred_linear is not None:
+                logits = self.graph_pred_linear(graph_rep)
+                return logits, node_rep
+            return graph_rep, node_rep
+
+        if self.graph_pred_linear is not None:
+            # If no batch given, default all nodes to batch 0
+            num_nodes = ops.shape(node_rep)[0]
+            batch_zeros = ops.zeros((num_nodes,), dtype="int32")
+            graph_rep = self.pool_fn(node_rep, batch_zeros)
+            logits = self.graph_pred_linear(graph_rep)
+            return logits, node_rep
+
+        return node_rep
+
+
+def load_mole_bert_weights(model: Union[MoleBERT, MoleBERTGNN], checkpoint_path: str):
+    """Loads official PyTorch Mole-BERT.pth checkpoint state dict into Keras 3 MoleBERT model."""
+    import torch
+
+    try:
+        state_dict = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
+    except Exception:
+        state_dict = torch.load(checkpoint_path, map_location="cpu")
+
+    if "state_dict" in state_dict:
+        state_dict = state_dict["state_dict"]
+
+    def get_np(key):
+        return state_dict[key].detach().cpu().float().numpy()
+
+    if not model.built:
+        model.build(None)
+
+    gnn = model.gnn if hasattr(model, "gnn") else model
+
+    # Embeddings
+    gnn.x_embedding1.set_weights([get_np("x_embedding1.weight")])
+    gnn.x_embedding2.set_weights([get_np("x_embedding2.weight")])
+
+    for l in range(gnn.num_layer):
+        conv = gnn.gnns[l]
+        bn = gnn.batch_norms[l]
+
+        # GINConv edge embeddings
+        conv.edge_embedding1.set_weights([get_np(f"gnns.{l}.edge_embedding1.weight")])
+        conv.edge_embedding2.set_weights([get_np(f"gnns.{l}.edge_embedding2.weight")])
+
+        # GINConv MLP layers
+        w0 = get_np(f"gnns.{l}.mlp.0.weight").T
+        b0 = get_np(f"gnns.{l}.mlp.0.bias")
+        conv.mlp.layers[0].set_weights([w0, b0])
+
+        w2 = get_np(f"gnns.{l}.mlp.2.weight").T
+        b2 = get_np(f"gnns.{l}.mlp.2.bias")
+        conv.mlp.layers[1].set_weights([w2, b2])
+
+        # Batch Normalization
+        gamma = get_np(f"batch_norms.{l}.weight")
+        beta = get_np(f"batch_norms.{l}.bias")
+        mean = get_np(f"batch_norms.{l}.running_mean")
+        var = get_np(f"batch_norms.{l}.running_var")
+        bn.set_weights([gamma, beta, mean, var])
+
+
+def download_mole_bert_checkpoint(cache_dir: Optional[str] = None) -> str:
+    """Downloads official Mole-BERT.pth checkpoint from GitHub."""
+    if cache_dir is None:
+        cache_dir = osp.expanduser("~/.cache/k3_node/mole_bert")
+
+    os.makedirs(cache_dir, exist_ok=True)
+    target_path = osp.join(cache_dir, "Mole-BERT.pth")
+
+    if osp.exists(target_path) and osp.getsize(target_path) > 1000:
+        return target_path
+
+    # Check local path in Mole-BERT/model_gin/Mole-BERT.pth
+    local_path = osp.join(
+        osp.dirname(osp.dirname(osp.dirname(osp.abspath(__file__)))),
+        "Mole-BERT",
+        "model_gin",
+        "Mole-BERT.pth",
+    )
+    if osp.exists(local_path) and osp.getsize(local_path) > 1000:
+        shutil.copyfile(local_path, target_path)
+        return target_path
+
+    return download_url(MOLE_BERT_URL, cache_dir, filename="Mole-BERT.pth")
+

k3_node/models/neural_fingerprint.py

@@ -0,0 +1,95 @@
+from typing import Optional
+
+import keras
+from keras import ops
+
+from k3_node.layers.conv import MFConv
+from k3_node.layers.pool import global_add_pool
+
+
+class NeuralFingerprint(keras.layers.Layer):
+    r"""The Neural Fingerprint model from the
+    `"Convolutional Networks on Graphs for Learning Molecular Fingerprints"
+    <https://arxiv.org/abs/1509.09292>`__ paper to generate fingerprints
+    of molecules.
+
+    Args:
+        in_channels (int): Size of each input sample.
+        hidden_channels (int): Size of each hidden sample.
+        out_channels (int): Size of each output fingerprint.
+        num_layers (int): Number of layers.
+        **kwargs (optional): Additional arguments of
+            :class:`~k3_node.layers.conv.MFConv`.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.models import NeuralFingerprint
+
+        x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
+        edge_index = np.random.randint(0, 10, size=(2, 30))  # 30 random edges
+
+        batch = np.repeat([0, 1], 5)  # two molecules with 5 atoms each
+        model = NeuralFingerprint(in_channels=8, hidden_channels=32, out_channels=16, num_layers=3)
+        fingerprint = model(x, edge_index, batch)  # one learned fingerprint per molecule
+        print(tuple(fingerprint.shape))  # (2, 16)
+        ```
+    """
+    def __init__(
+        self,
+        in_channels: int,
+        hidden_channels: int,
+        out_channels: int,
+        num_layers: int,
+        **kwargs,
+    ):
+        super().__init__()
+
+        self.in_channels = in_channels
+        self.hidden_channels = hidden_channels
+        self.out_channels = out_channels
+        self.num_layers = num_layers
+
+        self.convs = []
+        for i in range(self.num_layers):
+            in_c = self.in_channels if i == 0 else self.hidden_channels
+            self.convs.append(MFConv(in_c, hidden_channels, **kwargs))
+
+        self.lins = []
+        for _ in range(self.num_layers):
+            self.lins.append(keras.layers.Dense(out_channels, use_bias=False))
+
+    def build(self, input_shape=None):
+        self.built = True
+
+    def reset_parameters(self):
+        r"""Resets all learnable parameters of the module."""
+        for conv in self.convs:
+            if hasattr(conv, "reset_parameters"):
+                conv.reset_parameters()
+        for lin in self.lins:
+            if lin.built:
+                lin.kernel.assign(keras.initializers.GlorotUniform()(lin.kernel.shape))
+
+    def call(
+        self,
+        x,
+        edge_index,
+        batch: Optional[any] = None,
+        batch_size: Optional[int] = None,
+    ):
+        outs = []
+        for conv, lin in zip(self.convs, self.lins):
+            x = ops.sigmoid(conv(x, edge_index))
+            y = ops.softmax(lin(x), axis=-1)
+            outs.append(global_add_pool(y, batch, size=batch_size))
+
+        out = outs[0]
+        for item in outs[1:]:
+            out = out + item
+        return out
+
+    def __repr__(self) -> str:
+        return (f'{self.__class__.__name__}({self.in_channels}, '
+                f'{self.out_channels}, num_layers={self.num_layers})')
+

k3_node/models/node2vec.py

@@ -0,0 +1,213 @@
+from typing import Optional
+import numpy as np
+import keras
+from keras import ops
+
+
+class Node2Vec(keras.layers.Layer):
+    r"""The Node2Vec model from the
+    `"node2vec: Scalable Feature Learning for Networks"
+    <https://arxiv.org/abs/1607.00653>`_ paper where random walks of
+    length :obj:`walk_length` are sampled in a given graph, and node embeddings
+    are learned via negative sampling optimization.
+
+    Args:
+        edge_index: The edge indices.
+        embedding_dim (int): The size of each embedding vector.
+        walk_length (int): The walk length.
+        context_size (int): The actual context size which is considered for
+            positive samples.
+        walks_per_node (int, optional): The number of walks to sample for each
+            node. (default: :obj:`1`)
+        p (float, optional): Likelihood of immediately revisiting a node in the
+            walk. (default: :obj:`1.0`)
+        q (float, optional): Control parameter to interpolate between
+            breadth-first strategy and depth-first strategy. (default: :obj:`1.0`)
+        num_negative_samples (int, optional): The number of negative samples to
+            use for each positive sample. (default: :obj:`1`)
+        num_nodes (int, optional): The number of nodes. (default: :obj:`None`)
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.models import Node2Vec
+
+        edge_index = np.array([[0, 1, 2, 3, 0, 2], [1, 2, 3, 0, 2, 0]])
+        model = Node2Vec(edge_index, embedding_dim=16, walk_length=4, context_size=3, walks_per_node=2)
+        print(tuple(model().shape))  # (4, 16): embeddings of all nodes
+
+        batch = np.array([0, 1])
+        loss = model.loss(model.pos_sample(batch), model.neg_sample(batch))  # skip-gram loss on random walks
+        print(tuple(loss.shape))  # (): a scalar
+        ```
+    """
+    def __init__(
+        self,
+        edge_index,
+        embedding_dim: int,
+        walk_length: int,
+        context_size: int,
+        walks_per_node: int = 1,
+        p: float = 1.0,
+        q: float = 1.0,
+        num_negative_samples: int = 1,
+        num_nodes: Optional[int] = None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        edge_index_np = ops.convert_to_numpy(edge_index).astype(np.int64)
+        if num_nodes is None:
+            num_nodes = int(edge_index_np.max()) + 1 if edge_index_np.size > 0 else 0
+
+        self.num_nodes = num_nodes
+        self.embedding_dim = embedding_dim
+        self.walk_length = walk_length - 1
+        self.context_size = context_size
+        self.walks_per_node = walks_per_node
+        self.p = p
+        self.q = q
+        self.num_negative_samples = num_negative_samples
+        self.EPS = 1e-15
+
+        # Build adjacency list for fast random walks
+        self.adj = [[] for _ in range(num_nodes)]
+        if edge_index_np.size > 0:
+            for src, dst in zip(edge_index_np[0], edge_index_np[1]):
+                self.adj[src].append(int(dst))
+        self.adj_sets = [set(nbrs) for nbrs in self.adj]
+
+        self.embedding = keras.layers.Embedding(num_nodes, embedding_dim)
+
+    def reset_parameters(self):
+        if self.embedding.built:
+            self.embedding.embeddings.assign(
+                keras.initializers.GlorotUniform()(self.embedding.embeddings.shape)
+            )
+
+    def call(self, batch: Optional[any] = None):
+        """Returns the embeddings for the nodes in :obj:`batch`."""
+        if batch is None:
+            batch = ops.arange(self.num_nodes, dtype="int64")
+        return self.embedding(batch)
+
+    def pos_sample(self, batch):
+        batch_np = ops.convert_to_numpy(batch).astype(np.int64)
+        repeated = np.repeat(batch_np, self.walks_per_node)
+
+        # Sample random walks
+        all_walks = [self._random_walk(int(node)) for node in repeated]
+
+        rw = np.array(all_walks, dtype=np.int64)
+        walks = []
+        num_walks_per_rw = 1 + self.walk_length + 1 - self.context_size
+        for j in range(num_walks_per_rw):
+            walks.append(rw[:, j : j + self.context_size])
+        out = np.concatenate(walks, axis=0) if len(walks) > 0 else rw
+        return ops.convert_to_tensor(out, dtype="int64")
+
+    def _random_walk(self, start: int):
+        """A node2vec walk: from ``cur`` (reached from ``prev``) the next node is weighted by 1/p
+        if it returns to ``prev``, 1 if it is also a neighbor of ``prev`` and 1/q otherwise."""
+        walk = [start]
+        for _ in range(self.walk_length):
+            cur = walk[-1]
+            nbrs = self.adj[cur]
+            if not nbrs:
+                walk.append(cur)
+                continue
+            if (self.p == 1.0 and self.q == 1.0) or len(walk) == 1:
+                walk.append(nbrs[np.random.randint(len(nbrs))])
+                continue
+            prev = walk[-2]
+            prev_nbrs = self.adj_sets[prev]
+            weights = np.array([1.0 / self.p if n == prev else (1.0 if n in prev_nbrs else 1.0 / self.q)
+                                for n in nbrs])
+            walk.append(nbrs[np.random.choice(len(nbrs), p=weights / weights.sum())])
+        return walk
+
+    def loader(self, batch_size: int = 128, shuffle: bool = True):
+        r"""Yields ``(pos_rw, neg_rw)`` random-walk batches, starting from ``batch_size`` nodes
+        at a time, as PyG's ``Node2Vec.loader``."""
+        nodes = np.random.permutation(self.num_nodes) if shuffle else np.arange(self.num_nodes)
+        for start in range(0, self.num_nodes, batch_size):
+            batch = nodes[start:start + batch_size]
+            yield self.pos_sample(batch), self.neg_sample(batch)
+
+    def compile(self, optimizer):
+        r"""Sets the optimizer used by :meth:`fit`."""
+        self.optimizer = optimizer
+
+    def fit(self, epochs: int = 1, batch_size: int = 128, verbose: int = 1):
+        r"""Trains the embeddings on freshly sampled random walks for ``epochs`` passes over
+        all nodes; returns the mean loss of every epoch."""
+        from k3_node.training import gradient_step
+
+        if getattr(self, "optimizer", None) is None:
+            raise ValueError("Call `compile(optimizer=...)` before `fit`.")
+        self(ops.arange(1))  # create the embeddings
+        history = {"loss": []}
+        for epoch in range(1, epochs + 1):
+            losses = [gradient_step(lambda: self.loss(pos_rw, neg_rw), self.trainable_variables, self.optimizer)
+                      for pos_rw, neg_rw in self.loader(batch_size, shuffle=True)]
+            history["loss"].append(float(np.mean(losses)))
+            if verbose:
+                print(f"Epoch {epoch:03d}: loss: {history['loss'][-1]:.4f}")
+        return history
+
+    def test(self, train_z, train_y, test_z, test_y, solver: str = "lbfgs", *args, **kwargs):
+        r"""Evaluates the embeddings with a logistic regression classifier; returns its accuracy."""
+        from sklearn.linear_model import LogisticRegression
+
+        clf = LogisticRegression(*args, solver=solver, **kwargs).fit(
+            ops.convert_to_numpy(train_z), ops.convert_to_numpy(train_y))
+        return clf.score(ops.convert_to_numpy(test_z), ops.convert_to_numpy(test_y))
+
+    def neg_sample(self, batch):
+        batch_np = ops.convert_to_numpy(batch).astype(np.int64)
+        repeated = np.repeat(
+            batch_np, self.walks_per_node * self.num_negative_samples
+        )
+        rand_rest = np.random.randint(
+            0, max(self.num_nodes, 1), size=(len(repeated), self.walk_length)
+        )
+        rw = np.concatenate([repeated[:, None], rand_rest], axis=-1)
+
+        walks = []
+        num_walks_per_rw = 1 + self.walk_length + 1 - self.context_size
+        for j in range(num_walks_per_rw):
+            walks.append(rw[:, j : j + self.context_size])
+        out = np.concatenate(walks, axis=0) if len(walks) > 0 else rw
+        return ops.convert_to_tensor(out, dtype="int64")
+
+    def loss(self, pos_rw, neg_rw):
+        r"""Computes the loss given positive and negative random walks."""
+        # Positive loss
+        start = pos_rw[:, 0]
+        rest = pos_rw[:, 1:]
+        pos_b = ops.shape(pos_rw)[0]
+
+        h_start = ops.reshape(self.embedding(start), (pos_b, 1, self.embedding_dim))
+        h_rest = ops.reshape(
+            self.embedding(ops.reshape(rest, (-1,))),
+            (pos_b, -1, self.embedding_dim),
+        )
+
+        out = ops.reshape(ops.sum(h_start * h_rest, axis=-1), (-1,))
+        pos_loss = -ops.mean(ops.log(ops.sigmoid(out) + self.EPS))
+
+        # Negative loss
+        start = neg_rw[:, 0]
+        rest = neg_rw[:, 1:]
+        neg_b = ops.shape(neg_rw)[0]
+
+        h_start = ops.reshape(self.embedding(start), (neg_b, 1, self.embedding_dim))
+        h_rest = ops.reshape(
+            self.embedding(ops.reshape(rest, (-1,))),
+            (neg_b, -1, self.embedding_dim),
+        )
+
+        out = ops.reshape(ops.sum(h_start * h_rest, axis=-1), (-1,))
+        neg_loss = -ops.mean(ops.log(1.0 - ops.sigmoid(out) + self.EPS))
+
+        return pos_loss + neg_loss