k3-node 1.0.0__py3-none-any.whl

k3_node/ops/__init__.py

@@ -0,0 +1,4 @@
+from .graph import *
+from .matmul import *
+from .conv import *
+from .numpy import *

k3_node/ops/conv.py

@@ -0,0 +1,56 @@
+# ported from spektral
+
+import numpy as np
+import warnings, copy
+
+
+def degree_matrix(A):
+    degrees = np.array(A.sum(1)).flatten()
+    D = np.diag(degrees)
+    return D
+
+
+def degree_power(A, k):
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore")
+        degrees = np.power(np.array(A.sum(1)), k).ravel()
+    degrees[np.isinf(degrees)] = 0.0
+    D = np.diag(degrees)
+    return D
+
+
+def normalized_adjacency(A, symmetric=True):
+    if symmetric:
+        normalized_D = degree_power(A, -0.5)
+        return normalized_D.dot(A).dot(normalized_D)
+    else:
+        normalized_D = degree_power(A, -1.0)
+        return normalized_D.dot(A)
+
+
+def laplacian(A):
+    return degree_matrix(A) - A
+
+
+def normalized_laplacian(A, symmetric=True):
+    I = np.eye(A.shape[-1], dtype=A.dtype)
+    normalized_adj = normalized_adjacency(A, symmetric=symmetric)
+    return I - normalized_adj
+
+
+def gcn_filter(A, symmetric=True):
+    out = copy.deepcopy(A)
+    if isinstance(A, list) or (isinstance(A, np.ndarray) and A.ndim == 3):
+        for i in range(len(A)):
+            out[i] = A[i]
+            out[i][np.diag_indices_from(out[i])] += 1
+            out[i] = normalized_adjacency(out[i], symmetric=symmetric)
+    else:
+        if hasattr(out, "tocsr"):
+            out = out.tocsr()
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore")
+            out[np.diag_indices_from(out)] += 1
+        out = normalized_adjacency(out, symmetric=symmetric)
+
+    return out

k3_node/ops/creation.py

@@ -0,0 +1,43 @@
+"""Filled-tensor creation that also works on torch's ``meta`` device.
+
+``keras.ops.full`` / ``keras.ops.full_like`` call ``Tensor.item()`` on the torch backend, which
+fails on the ``meta`` device Keras uses to infer shapes before the first training step (the
+fallback then prints a long "unbuilt state" warning). ``ones * value`` avoids that.
+"""
+from keras import ops
+
+
+def full(shape, fill_value, dtype=None):
+    r"""Like :func:`keras.ops.full`."""
+    return ops.ones(shape, dtype=dtype) * ops.cast(fill_value, dtype or "float32")
+
+
+def full_like(x, fill_value, dtype=None):
+    r"""Like :func:`keras.ops.full_like`."""
+    return ops.ones_like(x, dtype=dtype) * ops.cast(fill_value, dtype or x.dtype)
+
+
+def scatter(indices, values, shape):
+    r"""Like :func:`keras.ops.scatter` (values at duplicate indices are summed)."""
+    from keras import backend
+
+    if backend.backend() == "torch":
+        import torch
+
+        values = values if torch.is_tensor(values) else ops.convert_to_tensor(values)
+        indices = indices if torch.is_tensor(indices) else ops.convert_to_tensor(indices)
+        out = torch.zeros(tuple(int(s) for s in shape), dtype=values.dtype, device=values.device)
+        return out.index_put_(tuple(indices.long().T), values, accumulate=True)
+    return ops.scatter(indices, values, shape)
+
+
+def repeat(x, repeats, axis=None):
+    r"""Like :func:`keras.ops.repeat`; an integer ``repeats`` also works on torch's ``meta`` device."""
+    from keras import backend
+
+    if backend.backend() == "torch" and isinstance(repeats, int):
+        import torch
+
+        x = x if torch.is_tensor(x) else ops.convert_to_tensor(x)
+        return torch.repeat_interleave(x.reshape(-1) if axis is None else x, repeats, dim=axis)
+    return ops.repeat(x, repeats, axis=axis)

k3_node/ops/graph.py

@@ -0,0 +1,27 @@
+# ported from spektral
+from keras import ops as k_ops, backend
+
+from k3_node.utils import *
+
+
+def degrees(A):
+    return k_ops.sum(A, axis=-1)
+
+
+def normalize_A(A):
+    D = degrees(A)
+    D = k_ops.sqrt(D)[:, None] + backend.epsilon()
+    perm = (0, 2, 1) if len(k_ops.shape(A)) == 3 else (1, 0)
+    output = (A / D) / k_ops.transpose(D, perm=perm)
+    return output
+
+
+def get_source_target(a):
+    if backend.backend() == "tensorflow":
+        import tensorflow as tf
+        if isinstance(a, tf.sparse.SparseTensor):
+            return a.indices[:, 0], a.indices[:, 1]
+        else:
+            return k_ops.where(a != 0)
+    else:
+        return k_ops.where(a != 0)

k3_node/ops/host.py

@@ -0,0 +1,41 @@
+import numpy as np
+from keras import ops
+
+
+def _is_meta(x) -> bool:
+    return getattr(getattr(x, "device", None), "type", None) == "meta"
+
+
+def to_numpy(x):
+    """``keras.ops.convert_to_numpy`` for code that runs on the host (NumPy).
+
+    During Keras' shape inference on the torch backend, tensors live on the "meta" device and have
+    no values. Host code then gets zeros of the right shape and dtype: only the output shapes matter
+    there, not the values.
+    """
+    if x is None or isinstance(x, np.ndarray):
+        return x
+    if _is_meta(x):
+        import torch
+
+        return np.zeros(tuple(x.shape), dtype=torch.empty(0, dtype=x.dtype).numpy().dtype)
+    try:
+        return ops.convert_to_numpy(x)
+    except Exception:
+        # TensorFlow and JAX trace symbolic tensors during Keras' shape inference. Unknown sizes
+        # (e.g. the number of nodes) become 1: zeros then describe a valid one-node graph.
+        if _in_shape_inference():
+            import keras
+
+            shape = tuple(1 if d is None else d for d in x.shape)
+            return np.zeros(shape, dtype=keras.backend.standardize_dtype(x.dtype))
+        raise
+
+
+def _in_shape_inference() -> bool:
+    try:
+        from keras.src.backend.common.symbolic_scope import in_symbolic_scope
+
+        return in_symbolic_scope()
+    except ImportError:
+        return False

k3_node/ops/matmul.py

@@ -0,0 +1,49 @@
+# ported from spektral
+
+from keras import ops as k_ops
+
+
+def dot(a, b):
+    a_ndim = len(k_ops.shape(a))
+    b_ndim = len(k_ops.shape(b))
+    assert a_ndim == b_ndim, "Expected equal ranks, got {} and {}" "".format(
+        a_ndim, b_ndim
+    )
+    return k_ops.matmul(a, b)
+
+
+def mixed_mode_dot(a, b):
+    a_shp = k_ops.shape(a)
+    b_shp = k_ops.shape(b)
+
+    b_t = k_ops.transpose(b, (1, 2, 0))
+    b_t = k_ops.reshape(b_t, k_ops.stack((b_shp[1], -1)))
+    output = dot(a, b_t)
+    output = k_ops.reshape(output, k_ops.stack((a_shp[0], b_shp[2], -1)))
+    output = k_ops.transpose(output, (2, 0, 1))
+
+    return output
+
+
+def modal_dot(a, b, transpose_a=False, transpose_b=False):
+    a_ndim = len(k_ops.shape(a))
+    b_ndim = len(k_ops.shape(b))
+    assert a_ndim in (2, 3), "Expected a of rank 2 or 3, got {}".format(a_ndim)
+    assert b_ndim in (2, 3), "Expected b of rank 2 or 3, got {}".format(b_ndim)
+
+    if transpose_a:
+        perm = (1, 0) if a_ndim == 2 else (0, 2, 1)
+        a = k_ops.transpose(a, perm)
+    if transpose_b:
+        perm = (1, 0) if b_ndim == 2 else (0, 2, 1)
+        b = k_ops.transpose(b, perm)
+    if a_ndim == b_ndim:
+        # ...ij,...jk->...ik
+        return dot(a, b)
+    elif a_ndim == 2:
+        # ij,bjk->bik
+        return mixed_mode_dot(a, b)
+    else:  # a_ndim == 3
+        # bij,jk->bik
+        # Immediately fallback to standard dense matmul, no need to reshape
+        return k_ops.matmul(a, b)

k3_node/ops/numpy.py

@@ -0,0 +1,24 @@
+from keras import backend, ops
+from k3_node.utils.backend_import import *
+
+
+def polyval(p, x):
+    p = ops.convert_to_tensor(p)
+
+    result = ops.zeros_like(x)
+
+    for i in range(p.shape[0]):
+        result = result * x + p[i]
+
+    return result
+
+
+def get_unique(inputs):
+    if backend.backend() == "tensorflow":
+        return tf.unique(inputs)
+    elif backend.backend() == "torch":
+        return torch.unique(inputs, return_inverse=True)
+    elif backend.backend() == "jax":
+        return jnp.unique(inputs, return_inverse=True)
+    elif backend.backend() == "numpy":
+        return np.unique(inputs, return_inverse=True)

k3_node/ops/segment.py

@@ -0,0 +1,54 @@
+"""Segment reductions (``segment_sum`` / ``segment_max`` / ``segment_min`` / ``segment_prod``).
+
+These match ``keras.ops.segment_*``. On the torch backend they are reimplemented so they also
+work on the ``meta`` device, which Keras uses to infer output shapes before the first training
+step: ``keras.ops.segment_*`` fails there (it repeats indices by a tensor-valued count), and the
+fallback it triggers runs on uninitialized memory, producing spurious "index out of range"
+warnings and occasionally huge allocations.
+"""
+from keras import backend, ops
+
+
+def _torch_segment(data, segment_ids, reduction, num_segments):
+    import torch
+
+    def as_tensor(x):
+        # keras.ops.convert_to_tensor would move tensors to the default device (a meta tensor
+        # becomes uninitialized memory), so keep existing tensors where they are.
+        if hasattr(x, "value") and not torch.is_tensor(x):
+            x = x.value
+        return x if torch.is_tensor(x) else ops.convert_to_tensor(x)
+
+    data, segment_ids = as_tensor(data), as_tensor(segment_ids).long()
+    if num_segments is None:
+        num_segments = int(segment_ids.max()) + 1 if segment_ids.numel() > 0 else 0
+    num_segments = int(num_segments)
+
+    # Out-of-range ids go to an extra segment that is dropped at the end (as in Keras).
+    segment_ids = torch.where((segment_ids >= 0) & (segment_ids < num_segments), segment_ids, num_segments)
+    index = segment_ids.reshape((-1,) + (1,) * (data.dim() - 1)).expand(data.shape)
+    fill = {"sum": 0.0, "amax": float("-inf"), "amin": float("inf"), "prod": 1.0}[reduction]
+    result = torch.full((num_segments + 1,) + tuple(data.shape[1:]), fill, device=data.device)
+    result = result.scatter_reduce(0, index, data.float(), reduction)
+    return result[:-1].to(data.dtype)
+
+
+def segment_sum(data, segment_ids, num_segments=None, sorted=False):
+    r"""Sums ``data`` over segments; like :func:`keras.ops.segment_sum`."""
+    if backend.backend() == "torch":
+        return _torch_segment(data, segment_ids, "sum", num_segments)
+    return ops.segment_sum(data, segment_ids, num_segments=num_segments, sorted=sorted)
+
+
+def segment_max(data, segment_ids, num_segments=None, sorted=False):
+    r"""Maximum of ``data`` over segments; like :func:`keras.ops.segment_max`."""
+    if backend.backend() == "torch":
+        return _torch_segment(data, segment_ids, "amax", num_segments)
+    return ops.segment_max(data, segment_ids, num_segments=num_segments, sorted=sorted)
+
+
+def segment_min(data, segment_ids, num_segments=None, sorted=False):
+    r"""Minimum of ``data`` over segments; like ``keras.ops.segment_min``."""
+    if backend.backend() == "torch":
+        return _torch_segment(data, segment_ids, "amin", num_segments)
+    return ops.segment_min(data, segment_ids, num_segments=num_segments, sorted=sorted)

k3_node/ops/sparse.py

@@ -0,0 +1,51 @@
+import keras
+from keras import ops
+
+from k3_node.ops.segment import segment_sum
+
+
+def spmm(index_targets, index_sources, edge_weight, x, num_targets: int):
+    r"""Sums weighted source features into targets:
+    :math:`\mathbf{out}_i = \sum_{(j, i)} w_{ji} \, \mathbf{x}_j`, i.e. a sparse-dense matrix product.
+
+    On torch this uses ``torch.sparse.mm``, which never materializes the ``[num_edges, channels]``
+    per-edge messages (the dominant memory cost of wide GCN-style layers). Other backends gather and
+    sum, which XLA fuses when compiled.
+
+    Args:
+        index_targets: Target node of every edge, shape ``[num_edges]``.
+        index_sources: Source node of every edge, shape ``[num_edges]``.
+        edge_weight: Edge weights of shape ``[num_edges]``, or ``None`` for weight 1.
+        x: Source node features of shape ``[num_sources, channels]``.
+        num_targets (int): The number of target nodes.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.ops.sparse import spmm
+
+        x = np.array([[1.0], [2.0], [3.0]], dtype="float32")  # one feature per node
+        targets, sources = np.array([0, 0, 2]), np.array([1, 2, 0])  # edges 1->0, 2->0, 0->2
+        weight = np.array([1.0, 0.5, 2.0], dtype="float32")
+
+        print(np.asarray(spmm(targets, sources, weight, x, 3)).ravel())  # [3.5 0.  2. ]
+        ```
+    """
+    if keras.config.backend() == "torch":
+        import torch
+
+        x_t = ops.convert_to_tensor(x)
+        if x_t.device.type != "meta" and len(x_t.shape) == 2:
+            index = torch.stack([ops.convert_to_tensor(index_targets).long(),
+                                 ops.convert_to_tensor(index_sources).long()])
+            if edge_weight is None:
+                values = torch.ones(index.shape[1], dtype=x_t.dtype, device=x_t.device)
+            else:
+                values = ops.convert_to_tensor(edge_weight).to(x_t.dtype).reshape(-1)
+            adj = torch.sparse_coo_tensor(index, values, (int(num_targets), x_t.shape[0]), check_invariants=True)
+            return torch.sparse.mm(adj, x_t)
+
+    messages = ops.take(x, index_sources, axis=0)
+    if edge_weight is not None:
+        messages = ops.expand_dims(ops.cast(edge_weight, messages.dtype), -1) * messages
+    return segment_sum(messages, ops.cast(index_targets, "int32"), num_segments=num_targets)

k3_node/rag/__init__.py

@@ -0,0 +1,49 @@
+"""GraphRAG & KG-LLM Connectors for K3-Node.
+
+Enables extracting subgraphs around retrieved entities, verbalizing structured knowledge
+into LLM prompt contexts (Llama 3, Mistral, ChatML), and projecting GNN (RGCN) and KGE (TransE)
+subgraph embeddings into dense prefix vectors for soft-prompt LLM augmentation.
+"""
+
+from k3_node.rag.subgraph import (
+    SubgraphResult,
+    extract_subgraph,
+    KGEntityRetriever,
+)
+from k3_node.rag.verbalizer import (
+    subgraph_to_triples,
+    verbalize_subgraph,
+    format_llm_prompt,
+)
+from k3_node.rag.encoders import (
+    RGCNSubGraphEncoder,
+    TransEPrefixEncoder,
+    GNNSubGraphEncoder,
+)
+from k3_node.rag.projector import (
+    GraphPrefixProjector,
+    KGLLMConnector,
+)
+from k3_node.rag.pipeline import (
+    GraphRAG,
+)
+
+__all__ = [
+    # Subgraph Extraction
+    "SubgraphResult",
+    "extract_subgraph",
+    "KGEntityRetriever",
+    # Verbalization & Prompts
+    "subgraph_to_triples",
+    "verbalize_subgraph",
+    "format_llm_prompt",
+    # GNN & KGE Encoders
+    "RGCNSubGraphEncoder",
+    "TransEPrefixEncoder",
+    "GNNSubGraphEncoder",
+    # Connectors & Projectors
+    "GraphPrefixProjector",
+    "KGLLMConnector",
+    # Unified Pipeline
+    "GraphRAG",
+]