k3-node 1.0.0__py3-none-any.whl

k3_node/test_docstring_examples.py

@@ -0,0 +1,106 @@
+"""Runs the ``Example:`` code in the docstrings of public layers and models.
+
+Every example must run on the current backend, and each ``print(...)  # expected`` line must print
+the value in its comment, so the documented output shapes cannot drift from the code.
+"""
+import ast
+import inspect
+import re
+import textwrap
+
+import keras
+import pytest
+
+import k3_node.layers
+import k3_node.models
+
+# Public objects that intentionally have no forward-pass example.
+NO_EXAMPLE = {
+    "Aggregation",  # abstract base class; see the concrete aggregations
+    "KGEModel",  # abstract base class; see TransE, DistMult, ComplEx, RotatE
+    "Connect",  # abstract base class of the pooling "connect" step
+    "Select",  # abstract base class of the pooling "select" step
+    "BasicGNN",  # abstract base class; see GCN, GraphSAGE, GIN, GAT, PNA, EdgeCNN
+}
+
+_FENCE = re.compile(r"```python\n(.*?)```", re.S)
+
+
+def _public_objects():
+    """Unique public layer/model classes (and layer functions), keyed by public name."""
+    found, seen = {}, set()
+    for package in (k3_node.layers, k3_node.models):
+        names = getattr(package, "__all__", None) or [n for n in dir(package) if not n.startswith("_")]
+        for name in sorted(names):
+            obj = getattr(package, name, None)
+            is_layer = inspect.isclass(obj) and issubclass(obj, keras.layers.Layer)
+            is_layer_fn = (
+                inspect.isfunction(obj) and package is k3_node.layers and obj.__module__.startswith("k3_node.layers")
+            )
+            if (is_layer or is_layer_fn) and id(obj) not in seen and obj.__module__.startswith("k3_node"):
+                seen.add(id(obj))
+                found[f"{package.__name__}.{name}"] = obj
+    return found
+
+
+def _other_documented_objects():
+    """Helpers outside of layers/models whose docstring examples are checked too."""
+    from k3_node import metrics, training, utils
+    from k3_node.data import Data
+    from k3_node.datasets import Digits, SEALDataset
+    from k3_node.ops.sparse import spmm
+    from k3_node.transforms import RandomLinkSplit
+
+    objects = [utils.normalized_cut, utils.k_hop_subgraph, utils.drnl_node_labeling, Digits, SEALDataset,
+               Data.edge_subgraph, RandomLinkSplit, spmm, training.gradient_step, metrics.F1Score]
+    return {f"{obj.__module__}.{obj.__qualname__}": obj for obj in objects}
+
+
+PUBLIC = {**_public_objects(), **_other_documented_objects()}
+
+
+def _examples(obj):
+    doc = obj.__doc__ or ""
+    section = doc[doc.find("Example") :] if "Example" in doc else ""
+    return [textwrap.dedent(block) for block in _FENCE.findall(section)]
+
+
+def _expected_prints(code):
+    """Maps the line number of each top-level print() to the text of its trailing comment."""
+    expected = {}
+    lines = code.split("\n")
+    for node in ast.parse(code).body:
+        if isinstance(node, ast.Expr) and isinstance(node.value, ast.Call) and getattr(node.value.func, "id", "") == "print":
+            _, sep, comment = lines[node.lineno - 1].partition("  # ")
+            if sep:
+                expected[node.lineno] = comment
+    return expected
+
+
+WITH_EXAMPLES = [name for name, obj in PUBLIC.items() if _examples(obj)]
+
+
+@pytest.mark.parametrize("name", WITH_EXAMPLES)
+def test_docstring_example_runs(name):
+    for code in _examples(PUBLIC[name]):
+        expected = _expected_prints(code)
+        printed = []
+
+        def capture(*args, **kwargs):
+            frame = inspect.currentframe().f_back
+            printed.append((frame.f_lineno, " ".join(str(a) for a in args)))
+
+        exec(compile(code, f"<{name} example>", "exec"), {"print": capture, "__name__": "__example__"})
+        for lineno, text in printed:
+            if lineno in expected:
+                assert expected[lineno].startswith(text), (
+                    f"{name}: line {lineno} printed {text!r}, but the docstring says {expected[lineno]!r}"
+                )
+
+
+def test_every_public_layer_and_model_has_an_example():
+    missing = sorted(
+        name for name, obj in PUBLIC.items()
+        if not _examples(obj) and name.rpartition(".")[2] not in NO_EXAMPLE and inspect.isclass(obj)
+    )
+    assert not missing, f"{len(missing)} public layers/models have no docstring example: {missing}"

k3_node/test_training_forwarding.py

@@ -0,0 +1,116 @@
+"""Every k3 layer must forward `training` to its dropout / batch-norm sublayers.
+
+On the JAX backend Keras does not propagate `training` from a model to nested layers during
+`fit`, so a layer that calls its dropout or batch norm without passing `training` silently runs
+it in inference mode. This test reproduces that: it disables Keras' propagation, enables dropout
+in every constructor, runs all docstring examples with `training=True`, and reports any nested
+call that did not receive `training` from its parent.
+"""
+import collections
+import inspect
+import re
+import textwrap
+
+import keras
+import pytest
+from keras.src.layers.layer import Layer
+
+import k3_node.layers
+import k3_node.models
+
+_TRAINING_SENSITIVE = ("BatchNormalization", "BatchNorm", "InstanceNorm", "HeteroBatchNorm")
+
+
+def _training_sensitive(layer):
+    for sub in layer._flatten_layers(include_self=True):
+        if type(sub).__name__ in _TRAINING_SENSITIVE:
+            return True
+        if isinstance(sub, keras.layers.Dropout) and float(sub.rate or 0) > 0:
+            return True
+        if isinstance(getattr(sub, "dropout", None), float) and sub.dropout > 0:  # e.g. attention dropout
+            return True
+    return False
+
+
+def _enable_dropout_defaults(cls):
+    """Sets zero-valued dropout defaults of cls.__init__ to 0.1 (so dropout paths are exercised)."""
+    init = cls.__dict__.get("__init__")
+    if init is None:
+        return None
+    saved = (init.__defaults__, dict(init.__kwdefaults__ or {}))
+    params = [p for p in inspect.signature(init).parameters.values() if p.kind in (p.POSITIONAL_ONLY, p.POSITIONAL_OR_KEYWORD)]
+    if init.__defaults__:
+        defaults = list(init.__defaults__)
+        for i, p in enumerate(params[len(params) - len(defaults):]):
+            if "drop" in p.name and not isinstance(defaults[i], bool) and defaults[i] == 0:
+                defaults[i] = 0.1
+        init.__defaults__ = tuple(defaults)
+    for k, v in (init.__kwdefaults__ or {}).items():
+        if "drop" in k and not isinstance(v, bool) and v == 0:
+            init.__kwdefaults__[k] = 0.1
+    return init, saved
+
+
+@pytest.fixture
+def no_training_propagation(monkeypatch):
+    stack, gaps = [], collections.Counter()
+    original_resolve = Layer._resolve_and_populate_arg
+    original_call = Layer.__call__
+
+    def resolve(self, arg_name, call_spec, call_context, kwargs):
+        if arg_name != "training":
+            return original_resolve(self, arg_name, call_spec, call_context, kwargs)
+        passed = arg_name in call_spec.user_arguments_dict
+        value = call_spec.user_arguments_dict.get(arg_name) if passed else (True if len(stack) <= 1 else None)
+        if self._call_has_context_arg.get(arg_name, False) and value is not None:
+            kwargs[arg_name] = value
+        parent = stack[-2] if len(stack) > 1 else None
+        if (parent is not None and not passed and type(parent).__module__.startswith("k3_node")
+                and self._call_has_context_arg.get("training", False) and _training_sensitive(self)):
+            gaps[f"{type(parent).__module__}.{type(parent).__name__} -> {type(self).__name__}"] += 1
+
+    def call(self, *args, **kwargs):
+        stack.append(self)
+        try:
+            return original_call(self, *args, **kwargs)
+        finally:
+            stack.pop()
+
+    monkeypatch.setattr(Layer, "_resolve_and_populate_arg", resolve)
+    monkeypatch.setattr(Layer, "__call__", call)
+    return gaps
+
+
+@pytest.mark.skipif(keras.backend.backend() != "torch", reason="checks the code, not a backend; torch is fastest")
+def test_layers_forward_training_to_dropout_and_batch_norm(no_training_propagation):
+    import sys
+
+    patched = []
+    for module in list(sys.modules.values()):
+        if getattr(module, "__name__", "").startswith("k3_node"):
+            for obj in list(vars(module).values()):
+                if inspect.isclass(obj) and issubclass(obj, Layer) and obj.__module__.startswith("k3_node"):
+                    result = _enable_dropout_defaults(obj)
+                    if result:
+                        patched.append(result)
+    try:
+        seen = set()
+        for package in (k3_node.layers, k3_node.models):
+            for name in dir(package):
+                obj = getattr(package, name)
+                if id(obj) in seen or not getattr(obj, "__module__", "").startswith("k3_node"):
+                    continue
+                seen.add(id(obj))
+                doc = obj.__doc__ or ""
+                for code in re.findall(r"```python\n(.*?)```", doc[doc.find("Example"):] if "Example" in doc else "", re.S):
+                    exec(compile(textwrap.dedent(code), name, "exec"), {"print": lambda *a, **k: None})
+    finally:
+        for init, (defaults, kwdefaults) in patched:
+            init.__defaults__ = defaults
+            if init.__kwdefaults__ is not None:
+                init.__kwdefaults__.clear()
+                init.__kwdefaults__.update(kwdefaults)
+    assert not no_training_propagation, (
+        "These layers call a dropout / batch-norm sublayer without forwarding `training` "
+        f"(it would never train on JAX): {sorted(no_training_propagation)}"
+    )

k3_node/training.py

@@ -0,0 +1,115 @@
+"""Backend-agnostic training helpers for models that don't fit ``keras.Model.fit``.
+
+Examples are models with several optimizers (like the adversarial autoencoders) or losses that
+are computed outside of a model's ``call``.
+"""
+from typing import Callable, Sequence
+
+import keras
+from keras import ops
+
+
+def _seed_state():
+    from keras.src.random.seed_generator import global_seed_generator
+
+    return global_seed_generator().state
+
+
+def _scalar(loss):
+    return ops.mean(loss) if len(ops.shape(loss)) > 0 else loss
+
+
+def no_grad():
+    r"""A context manager for evaluating a model outside of ``fit`` / ``predict``, like PyTorch's
+    ``torch.no_grad()``.
+
+    On the torch backend, calling a model records every intermediate result for a possible
+    backward pass, which can take far more memory than the model itself. Inside ``no_grad()``
+    nothing is recorded. TensorFlow and JAX only record gradients when asked to, so there it does
+    nothing.
+
+    Example:
+        ```python
+        import keras
+        import numpy as np
+        from k3_node.training import no_grad
+
+        model = keras.layers.Dense(4)
+        with no_grad():
+            out = model(np.random.rand(10, 8).astype("float32"))
+        print(tuple(out.shape))  # (10, 4)
+        ```
+    """
+    if keras.config.backend() == "torch":
+        import torch
+
+        return torch.no_grad()
+    import contextlib
+
+    return contextlib.nullcontext()
+
+
+def gradient_step(loss_fn: Callable, variables: Sequence, optimizer, state_variables: Sequence = ()):
+    r"""Runs ``loss_fn()``, then updates ``variables`` with one ``optimizer`` step on its gradients.
+
+    Works eagerly on every backend. Variables used by ``loss_fn`` but not listed in ``variables``
+    are not trained. ``state_variables`` are non-trainable variables that ``loss_fn`` may update
+    (for example the moving statistics of batch normalization); this matters on JAX only.
+
+    Args:
+        loss_fn (callable): A function without arguments that returns the loss (a scalar, or
+            per-example losses, which are averaged).
+        variables (list): The variables to train, e.g. ``model.trainable_variables``.
+        optimizer (keras.optimizers.Optimizer): The optimizer applying the update.
+        state_variables (list, optional): Non-trainable variables updated by ``loss_fn``.
+
+    Returns:
+        The loss as a Python float.
+
+    Example:
+        ```python
+        import keras
+        from keras import ops
+        from k3_node.training import gradient_step
+
+        w = keras.Variable(3.0)
+        optimizer = keras.optimizers.SGD(learning_rate=0.25)
+        loss = gradient_step(lambda: ops.square(w), [w], optimizer)  # gradient 2 * w = 6
+        print(loss, float(ops.convert_to_numpy(w)))  # 9.0 1.5
+        ```
+    """
+    variables = list(variables)
+    backend = keras.config.backend()
+
+    if backend == "tensorflow":
+        import tensorflow as tf
+
+        with tf.GradientTape() as tape:
+            loss = _scalar(loss_fn())
+        grads = tape.gradient(loss, variables)
+    elif backend == "torch":
+        import torch
+
+        loss = _scalar(loss_fn())
+        grads = torch.autograd.grad(loss, [v.value for v in variables], allow_unused=True)
+    elif backend == "jax":
+        import jax
+        from keras.src.backend.common.stateless_scope import StatelessScope
+
+        tracked = list(state_variables) + [_seed_state()]
+
+        def compute(values):
+            with StatelessScope(state_mapping=list(zip(variables, values))) as scope:
+                out = _scalar(loss_fn())
+            return out, [scope.get_current_value(v) for v in tracked]
+
+        (loss, updates), grads = jax.value_and_grad(compute, has_aux=True)([v.value for v in variables])
+        for v, value in zip(tracked, updates):
+            if value is not None:
+                v.assign(value)
+    else:
+        raise NotImplementedError(f"gradient_step does not support the {backend} backend.")
+
+    grads = [ops.zeros_like(v) if g is None else g for g, v in zip(grads, variables)]
+    optimizer.apply_gradients(zip(grads, variables))
+    return float(ops.convert_to_numpy(loss))

k3_node/transforms/__init__.py

@@ -0,0 +1,166 @@
+from k3_node.transforms.base_transform import BaseTransform, functional_transform
+from k3_node.transforms.compose import Compose, ComposeFilters
+
+from k3_node.transforms.general import (
+    ToDevice,
+    ToSparseTensor,
+    Constant,
+    NormalizeFeatures,
+    SVDFeatureReduction,
+    RemoveTrainingClasses,
+    RandomNodeSplit,
+    RandomLinkSplit,
+    AttentiveFPFeatures,
+    CompleteGraph,
+    NodePropertySplit,
+    IndexToMask,
+    MaskToIndex,
+    Pad,
+    Padding,
+    UniformPadding,
+    MappingPadding,
+)
+
+from k3_node.transforms.graph import (
+    ToUndirected,
+    OneHotDegree,
+    TargetIndegree,
+    LocalDegreeProfile,
+    AddSelfLoops,
+    AddRemainingSelfLoops,
+    RemoveSelfLoops,
+    RemoveIsolatedNodes,
+    RemoveDuplicatedEdges,
+    KNNGraph,
+    RadiusGraph,
+    ToDense,
+    TwoHop,
+    LineGraph,
+    LaplacianLambdaMax,
+    GDC,
+    SIGN,
+    GCNNorm,
+    AddMetaPaths,
+    AddRandomMetaPaths,
+    RootedEgoNets,
+    RootedRWSubgraph,
+    LargestConnectedComponents,
+    VirtualNode,
+    AddLaplacianEigenvectorPE,
+    AddRandomWalkPE,
+    AddGPSE,
+    FeaturePropagation,
+    HalfHop,
+)
+
+from k3_node.transforms.spatial import (
+    Distance,
+    Cartesian,
+    LocalCartesian,
+    Polar,
+    Spherical,
+    PointPairFeatures,
+    Center,
+    NormalizeRotation,
+    NormalizeScale,
+    RandomJitter,
+    RandomFlip,
+    LinearTransformation,
+    RandomScale,
+    RandomRotate,
+    RandomShear,
+    FaceToEdge,
+    SamplePoints,
+    FixedPoints,
+    GenerateMeshNormals,
+    Delaunay,
+    ToSLIC,
+    GridSampling,
+    RandomTranslate,
+)
+
+general_transforms = [
+    'BaseTransform',
+    'Compose',
+    'ComposeFilters',
+    'ToDevice',
+    'ToSparseTensor',
+    'Constant',
+    'NormalizeFeatures',
+    'SVDFeatureReduction',
+    'RemoveTrainingClasses',
+    'RandomNodeSplit',
+    'RandomLinkSplit',
+    'AttentiveFPFeatures',
+    'CompleteGraph',
+    'NodePropertySplit',
+    'IndexToMask',
+    'MaskToIndex',
+    'Pad',
+]
+
+graph_transforms = [
+    'ToUndirected',
+    'OneHotDegree',
+    'TargetIndegree',
+    'LocalDegreeProfile',
+    'AddSelfLoops',
+    'AddRemainingSelfLoops',
+    'RemoveSelfLoops',
+    'RemoveIsolatedNodes',
+    'RemoveDuplicatedEdges',
+    'KNNGraph',
+    'RadiusGraph',
+    'ToDense',
+    'TwoHop',
+    'LineGraph',
+    'LaplacianLambdaMax',
+    'GDC',
+    'SIGN',
+    'GCNNorm',
+    'AddMetaPaths',
+    'AddRandomMetaPaths',
+    'RootedEgoNets',
+    'RootedRWSubgraph',
+    'LargestConnectedComponents',
+    'VirtualNode',
+    'AddLaplacianEigenvectorPE',
+    'AddRandomWalkPE',
+    'AddGPSE',
+    'FeaturePropagation',
+    'HalfHop',
+]
+
+vision_transforms = [
+    'Distance',
+    'Cartesian',
+    'LocalCartesian',
+    'Polar',
+    'Spherical',
+    'PointPairFeatures',
+    'Center',
+    'NormalizeRotation',
+    'NormalizeScale',
+    'RandomJitter',
+    'RandomFlip',
+    'LinearTransformation',
+    'RandomScale',
+    'RandomRotate',
+    'RandomShear',
+    'FaceToEdge',
+    'SamplePoints',
+    'FixedPoints',
+    'GenerateMeshNormals',
+    'Delaunay',
+    'ToSLIC',
+    'GridSampling',
+]
+
+__all__ = general_transforms + graph_transforms + vision_transforms + [
+    'RandomTranslate',
+    'Padding',
+    'UniformPadding',
+    'MappingPadding',
+    'functional_transform',
+]
+

k3_node/transforms/base_transform.py

@@ -0,0 +1,32 @@
+import copy
+from abc import ABC, abstractmethod
+from typing import Any, Callable
+
+
+class BaseTransform(ABC):
+    r"""An abstract base class for writing transforms.
+
+    Transforms are a general way to modify and customize
+    :class:`~k3_node.data.Data` or :class:`~k3_node.data.HeteroData` objects.
+    """
+
+    def __call__(self, data: Any) -> Any:
+        # Shallow-copy the data to prevent in-place modification of caller's object
+        return self.forward(copy.copy(data))
+
+    @abstractmethod
+    def forward(self, data: Any) -> Any:
+        pass
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}()"
+
+
+def functional_transform(name: str) -> Callable:
+    r"""Decorator for functional transforms."""
+
+    def wrapper(cls: Any) -> Any:
+        return cls
+
+    return wrapper
+

k3_node/transforms/compose.py

@@ -0,0 +1,58 @@
+from typing import Callable, List, Union
+
+from k3_node.data import Data, HeteroData
+from k3_node.transforms.base_transform import BaseTransform
+
+
+class Compose(BaseTransform):
+    r"""Composes several transforms together.
+
+    Args:
+        transforms (List[Callable]): List of transforms to compose.
+    """
+
+    def __init__(self, transforms: List[Callable]):
+        self.transforms = transforms
+
+    def forward(
+        self,
+        data: Union[Data, HeteroData],
+    ) -> Union[Data, HeteroData]:
+        for transform in self.transforms:
+            if isinstance(data, (list, tuple)):
+                data = [transform(d) for d in data]
+            else:
+                data = transform(data)
+        return data
+
+    def __repr__(self) -> str:
+        args = [f"  {transform}" for transform in self.transforms]
+        return "{}([\n{}\n])".format(self.__class__.__name__, ",\n".join(args))
+
+
+class ComposeFilters:
+    r"""Composes several filters together.
+
+    Args:
+        filters (List[Callable]): List of filters to compose.
+    """
+
+    def __init__(self, filters: List[Callable]):
+        self.filters = filters
+
+    def __call__(
+        self,
+        data: Union[Data, HeteroData],
+    ) -> bool:
+        for filter_fn in self.filters:
+            if isinstance(data, (list, tuple)):
+                if not all([filter_fn(d) for d in data]):
+                    return False
+            elif not filter_fn(data):
+                return False
+        return True
+
+    def __repr__(self) -> str:
+        args = [f"  {filter_fn}" for filter_fn in self.filters]
+        return "{}([\n{}\n])".format(self.__class__.__name__, ",\n".join(args))
+