k3-node 1.0.0__py3-none-any.whl

k3_node/models/basic_gnn.py

@@ -0,0 +1,443 @@
+import copy
+import inspect
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+
+import keras
+from keras import ops
+
+from k3_node.layers.conv import (
+    EdgeConv,
+    GATConv,
+    GATv2Conv,
+    GCNConv,
+    GINConv,
+    MessagePassing,
+    PNAConv,
+    SAGEConv,
+)
+from k3_node.models.mlp import MLP, _normalization_resolver
+from k3_node.models.jumping_knowledge import JumpingKnowledge
+from k3_node.hub.hub_mixin import K3NodeHubMixin
+
+
+class BasicGNN(K3NodeHubMixin, keras.Model):
+    r"""An abstract base class for implementing basic GNN models.
+
+    Args:
+        in_channels (int or tuple): Size of each input sample.
+        hidden_channels (int): Size of each hidden sample.
+        num_layers (int): Number of message passing layers.
+        out_channels (int, optional): If not set to :obj:`None`, will apply a
+            final linear transformation to convert hidden node embeddings to
+            output size :obj:`out_channels`. (default: :obj:`None`)
+        dropout (float, optional): Dropout probability. (default: :obj:`0.`)
+        act (str or Callable, optional): The non-linear activation function to
+            use. (default: :obj:`"relu"`)
+        act_first (bool, optional): If set to :obj:`True`, activation is
+            applied before normalization. (default: :obj:`False`)
+        act_kwargs (Dict[str, Any], optional): Arguments passed to the
+            respective activation function defined by :obj:`act`.
+            (default: :obj:`None`)
+        norm (str or Callable, optional): The normalization function to
+            use. (default: :obj:`None`)
+        norm_kwargs (Dict[str, Any], optional): Arguments passed to the
+            respective normalization function defined by :obj:`norm`.
+            (default: :obj:`None`)
+        jk (str, optional): The Jumping Knowledge mode. If specified, the model
+            will additionally apply a final linear transformation to transform
+            node embeddings to the expected output feature dimensionality.
+            (:obj:`None`, :obj:`"last"`, :obj:`"cat"`, :obj:`"max"`,
+            :obj:`"lstm"`). (default: :obj:`None`)
+        **kwargs (optional): Additional arguments of the underlying
+            :class:`torch_geometric.nn.conv.MessagePassing` layers.
+    
+    Call arguments: ``x``, ``edge_index`` and optionally ``edge_weight``, ``edge_attr`` and
+    ``batch``. With mini-batches from :class:`~k3_node.loader.NeighborLoader`, also pass
+    ``num_sampled_nodes_per_hop`` and ``num_sampled_edges_per_hop`` (the batch's
+    ``num_sampled_nodes`` / ``num_sampled_edges``) to skip the nodes and edges that no longer
+    affect the seed nodes after each layer, as in PyG. This trimming needs concrete sizes, so
+    compile the model with ``run_eagerly=True`` when using it.
+    """
+    supports_edge_weight: bool = False
+    supports_edge_attr: bool = False
+    supports_norm_batch: bool = False
+
+    def __init__(
+        self,
+        in_channels: int,
+        hidden_channels: int,
+        num_layers: int,
+        out_channels: Optional[int] = None,
+        dropout: float = 0.0,
+        act: Union[str, Callable, None] = "relu",
+        act_first: bool = False,
+        act_kwargs: Optional[Dict[str, Any]] = None,
+        norm: Union[str, Callable, None] = None,
+        norm_kwargs: Optional[Dict[str, Any]] = None,
+        jk: Optional[str] = None,
+        **kwargs,
+    ):
+        super().__init__()
+
+        self.in_channels = in_channels
+        self.hidden_channels = hidden_channels
+        self.num_layers = num_layers
+        dropout = float(dropout) if dropout is not None else 0.0
+        self.dropout_p = dropout
+        self.dropout = keras.layers.Dropout(rate=dropout) if dropout > 0 else None
+
+        if isinstance(act, str):
+            self.act = keras.activations.get(act)
+        else:
+            self.act = act
+
+        self.jk_mode = jk
+        self.act_first = act_first
+        self.norm_query = norm
+        self.norm_kwargs = norm_kwargs or {}
+
+        if out_channels is not None:
+            self.out_channels = out_channels
+        else:
+            self.out_channels = hidden_channels
+
+        self.convs = []
+        curr_in = in_channels
+        if num_layers > 1:
+            self.convs.append(self.init_conv(curr_in, hidden_channels, **kwargs))
+            if isinstance(curr_in, (tuple, list)):
+                curr_in = (hidden_channels, hidden_channels)
+            else:
+                curr_in = hidden_channels
+
+        for _ in range(num_layers - 2):
+            self.convs.append(self.init_conv(curr_in, hidden_channels, **kwargs))
+            if isinstance(curr_in, (tuple, list)):
+                curr_in = (hidden_channels, hidden_channels)
+            else:
+                curr_in = hidden_channels
+
+        if out_channels is not None and jk is None:
+            self._is_conv_to_out = True
+            self.convs.append(self.init_conv(curr_in, out_channels, **kwargs))
+        else:
+            self.convs.append(self.init_conv(curr_in, hidden_channels, **kwargs))
+
+        self.norms = []
+        self.supports_norm_batch = False
+
+        for _ in range(num_layers - 1):
+            if norm is not None:
+                norm_layer = _normalization_resolver(norm, hidden_channels, **self.norm_kwargs)
+                self.norms.append(norm_layer)
+                if hasattr(norm_layer, "call"):
+                    sig = inspect.signature(norm_layer.call).parameters
+                    self.supports_norm_batch = "batch" in sig
+            else:
+                self.norms.append(None)
+
+        if jk is not None:
+            if norm is not None:
+                self.norms.append(_normalization_resolver(norm, hidden_channels, **self.norm_kwargs))
+            else:
+                self.norms.append(None)
+        else:
+            self.norms.append(None)
+
+        if jk is not None and jk != "last":
+            self.jk = JumpingKnowledge(jk, hidden_channels, num_layers)
+
+        if jk is not None:
+            if jk == "cat":
+                jk_in = num_layers * hidden_channels
+            else:
+                jk_in = hidden_channels
+            self.lin = keras.layers.Dense(self.out_channels)
+
+    def init_conv(self, in_channels: Union[int, Tuple[int, int]],
+                  out_channels: int, **kwargs) -> MessagePassing:
+        raise NotImplementedError
+
+    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 norm in self.norms:
+            if norm is not None and hasattr(norm, "reset_parameters"):
+                norm.reset_parameters()
+        if hasattr(self, "jk") and hasattr(self.jk, "reset_parameters"):
+            self.jk.reset_parameters()
+        if hasattr(self, "lin") and hasattr(self.lin, "reset_parameters"):
+            self.lin.reset_parameters()
+
+    def call(
+        self,
+        x,
+        edge_index=None,
+        edge_weight=None,
+        edge_attr=None,
+        batch=None,
+        batch_size=None,
+        num_sampled_nodes_per_hop=None,
+        num_sampled_edges_per_hop=None,
+        training=None,
+    ):
+        if hasattr(x, "edge_index") and edge_index is None:
+            edge_index = getattr(x, "edge_index", None)
+            edge_weight = getattr(x, "edge_weight", None) if edge_weight is None else edge_weight
+            edge_attr = getattr(x, "edge_attr", None) if edge_attr is None else edge_attr
+            batch = getattr(x, "batch", None) if batch is None else batch
+            x = x.x
+        elif isinstance(x, (tuple, list)) and edge_index is None:
+            if len(x) > 1:
+                edge_index = x[1]
+            if len(x) > 2:
+                edge_attr = x[2]
+            x = x[0]
+
+        trim = num_sampled_nodes_per_hop is not None and num_sampled_edges_per_hop is not None
+        if trim:
+            from k3_node.ops.host import to_numpy
+
+            nodes_per_hop = [int(n) for n in to_numpy(num_sampled_nodes_per_hop)]
+            edges_per_hop = [int(n) for n in to_numpy(num_sampled_edges_per_hop)]
+
+        xs: List = []
+        # `training` is forwarded explicitly: Keras does not propagate it to nested layers on JAX.
+        for i, (conv, norm) in enumerate(zip(self.convs, self.norms)):
+            if trim and i > 0:
+                # Hierarchical neighborhood sampling: the nodes and edges of the outermost hop no
+                # longer influence the seed nodes, so drop them (as PyG's `trim_to_layer`).
+                x = x[: x.shape[0] - nodes_per_hop[-i]]
+                num_edges = edge_index.shape[1] - edges_per_hop[-i]
+                edge_index = edge_index[:, :num_edges]
+                edge_weight = None if edge_weight is None else edge_weight[:num_edges]
+                edge_attr = None if edge_attr is None else edge_attr[:num_edges]
+            if self.supports_edge_weight and self.supports_edge_attr:
+                x = conv(x, edge_index, edge_weight=edge_weight, edge_attr=edge_attr, training=training)
+            elif self.supports_edge_weight:
+                x = conv(x, edge_index, edge_weight=edge_weight, training=training)
+            elif self.supports_edge_attr:
+                x = conv(x, edge_index, edge_attr=edge_attr, training=training)
+            else:
+                x = conv(x, edge_index, training=training)
+
+            if i < self.num_layers - 1 or self.jk_mode is not None:
+                if self.act is not None and self.act_first:
+                    x = self.act(x)
+                if norm is not None:
+                    if self.supports_norm_batch and batch is not None:
+                        x = norm(x, batch=batch, training=training)
+                    else:
+                        x = norm(x, training=training)
+                if self.act is not None and not self.act_first:
+                    x = self.act(x)
+                if self.dropout is not None:
+                    x = self.dropout(x, training=training)
+                if hasattr(self, "jk"):
+                    xs.append(x)
+
+        if hasattr(self, "jk"):
+            x = self.jk(xs)
+        if hasattr(self, "lin"):
+            x = self.lin(x)
+
+        return x
+
+    def __repr__(self) -> str:
+        return (f'{self.__class__.__name__}({self.in_channels}, '
+                f'{self.out_channels}, num_layers={self.num_layers})')
+
+
+class GCN(BasicGNN):
+    r"""The Graph Neural Network from the `"Semi-supervised
+    Classification with Graph Convolutional Networks"
+    <https://arxiv.org/abs/1609.02907>`_ paper, using the
+    :class:`~k3_node.layers.conv.GCNConv` operator for message passing.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.models import GCN
+
+        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
+
+        model = GCN(in_channels=8, hidden_channels=16, num_layers=2, out_channels=4)
+        out = model(x, edge_index)  # e.g. logits for 4 classes per node
+        print(tuple(out.shape))  # (10, 4)
+        ```
+    """
+    supports_edge_weight: bool = True
+    supports_edge_attr: bool = False
+
+    def init_conv(self, in_channels: int, out_channels: int, **kwargs) -> MessagePassing:
+        return GCNConv(in_channels, out_channels, **kwargs)
+
+
+class GraphSAGE(BasicGNN):
+    r"""The Graph Neural Network from the `"Inductive Representation Learning
+    on Large Graphs" <https://arxiv.org/abs/1706.02216>`_ paper, using the
+    :class:`~k3_node.layers.conv.SAGEConv` operator for message passing.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.models import GraphSAGE
+
+        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
+
+        model = GraphSAGE(in_channels=8, hidden_channels=16, num_layers=2, out_channels=4)
+        out = model(x, edge_index)  # e.g. logits for 4 classes per node
+        print(tuple(out.shape))  # (10, 4)
+        ```
+    """
+    supports_edge_weight: bool = False
+    supports_edge_attr: bool = False
+
+    def init_conv(self, in_channels: Union[int, Tuple[int, int]],
+                  out_channels: int, **kwargs) -> MessagePassing:
+        return SAGEConv(in_channels, out_channels, **kwargs)
+
+
+class GIN(BasicGNN):
+    r"""The Graph Neural Network from the `"How Powerful are Graph Neural
+    Networks?" <https://arxiv.org/abs/1810.00826>`_ paper, using the
+    :class:`~k3_node.layers.conv.GINConv` operator for message passing.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.models import GIN
+
+        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
+
+        model = GIN(in_channels=8, hidden_channels=16, num_layers=2, out_channels=4)
+        out = model(x, edge_index)  # e.g. logits for 4 classes per node
+        print(tuple(out.shape))  # (10, 4)
+        ```
+    """
+    supports_edge_weight: bool = False
+    supports_edge_attr: bool = False
+
+    def init_conv(self, in_channels: int, out_channels: int, **kwargs) -> MessagePassing:
+        mlp = MLP(
+            [in_channels, out_channels, out_channels],
+            act=self.act,
+            act_first=self.act_first,
+            norm=self.norm_query,
+            norm_kwargs=self.norm_kwargs,
+        )
+        return GINConv(mlp, **kwargs)
+
+
+class GAT(BasicGNN):
+    r"""The Graph Neural Network from `"Graph Attention Networks"
+    <https://arxiv.org/abs/1710.10903>`_ or `"How Attentive are Graph Attention
+    Networks?" <https://arxiv.org/abs/2105.14491>`_ papers, using the
+    :class:`~k3_node.layers.conv.GATConv` or
+    :class:`~k3_node.layers.conv.GATv2Conv` operator for message passing.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.models import GAT
+
+        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
+
+        model = GAT(in_channels=8, hidden_channels=16, num_layers=2, out_channels=4, heads=2)
+        out = model(x, edge_index)  # e.g. logits for 4 classes per node
+        print(tuple(out.shape))  # (10, 4)
+        ```
+    """
+    supports_edge_weight: bool = False
+    supports_edge_attr: bool = True
+
+    def init_conv(self, in_channels: Union[int, Tuple[int, int]],
+                  out_channels: int, **kwargs) -> MessagePassing:
+        v2 = kwargs.pop('v2', False)
+        heads = kwargs.pop('heads', 1)
+        concat = kwargs.pop('concat', True)
+
+        if getattr(self, '_is_conv_to_out', False):
+            concat = False
+
+        if concat and out_channels % heads != 0:
+            raise ValueError(f"Ensure that the number of output channels of "
+                             f"'GATConv' (got '{out_channels}') is divisible "
+                             f"by the number of heads (got '{heads}')")
+
+        if concat:
+            out_channels = out_channels // heads
+
+        Conv = GATConv if not v2 else GATv2Conv
+        return Conv(in_channels, out_channels, heads=heads, concat=concat,
+                    dropout=self.dropout_p, **kwargs)
+
+
+class PNA(BasicGNN):
+    r"""The Graph Neural Network from the `"Principal Neighbourhood Aggregation
+    for Graph Nets" <https://arxiv.org/abs/2004.05718>`_ paper, using the
+    :class:`~k3_node.layers.conv.PNAConv` operator for message passing.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.models import PNA
+
+        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
+
+        deg = np.array([0, 2, 4, 3, 1])  # in-degree histogram of the training graphs
+        model = PNA(in_channels=8, hidden_channels=16, num_layers=2, out_channels=4,
+                    aggregators=["mean", "min", "max", "std"],
+                    scalers=["identity", "amplification", "attenuation"], deg=deg)
+        out = model(x, edge_index)
+        print(tuple(out.shape))  # (10, 4)
+        ```
+    """
+    supports_edge_weight: bool = False
+    supports_edge_attr: bool = True
+
+    def init_conv(self, in_channels: int, out_channels: int, **kwargs) -> MessagePassing:
+        return PNAConv(in_channels, out_channels, **kwargs)
+
+
+class EdgeCNN(BasicGNN):
+    r"""The Graph Neural Network from the `"Dynamic Graph CNN for Learning on
+    Point Clouds" <https://arxiv.org/abs/1801.07829>`_ paper, using the
+    :class:`~k3_node.layers.conv.EdgeConv` operator for message passing.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.models import EdgeCNN
+
+        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
+
+        model = EdgeCNN(in_channels=8, hidden_channels=16, num_layers=2, out_channels=4)
+        out = model(x, edge_index)  # e.g. logits for 4 classes per node
+        print(tuple(out.shape))  # (10, 4)
+        ```
+    """
+    supports_edge_weight: bool = False
+    supports_edge_attr: bool = False
+
+    def init_conv(self, in_channels: int, out_channels: int, **kwargs) -> MessagePassing:
+        mlp = MLP(
+            [2 * in_channels, out_channels, out_channels],
+            act=self.act,
+            act_first=self.act_first,
+            norm=self.norm_query,
+            norm_kwargs=self.norm_kwargs,
+        )
+        return EdgeConv(mlp, **kwargs)
+

k3_node/models/bio/__init__.py

@@ -0,0 +1,4 @@
+"""Biological and macromolecular graph neural network models (aliased from k3_node.applications.bio)."""
+
+from k3_node.applications.bio import *
+from k3_node.applications.bio import __all__

k3_node/models/captum.py

@@ -0,0 +1,52 @@
+r"""Captum model interpretability integration.
+
+Note: Captum is a PyTorch-specific model interpretability library.
+In k3-node (a multi-backend Keras 3 graph library supporting TensorFlow,
+PyTorch, and JAX), Captum integrations are provided for PyTorch backend
+compatibility or documented as PyTorch-exclusive.
+"""
+
+from typing import Optional, Union, Any
+
+
+def to_captum_model(
+    model: Any,
+    mask_type: str = "edge",
+    output_idx: Optional[int] = None,
+    metadata: Optional[Any] = None,
+):
+    r"""Converts a model into a Captum-compatible module.
+
+    .. note::
+        Captum is a PyTorch-exclusive library. This function requires
+        PyTorch backend and torch.nn.Module models.
+    """
+    raise NotImplementedError(
+        "Captum integration is PyTorch-specific and requires native torch.nn.Module."
+    )
+
+
+def to_captum_input(
+    x: Any,
+    edge_index: Any,
+    mask_type: str = "edge",
+    *args,
+    **kwargs,
+):
+    r"""Converts graph inputs into Captum-compatible inputs."""
+    raise NotImplementedError(
+        "Captum integration is PyTorch-specific and requires native torch.nn.Module."
+    )
+
+
+def captum_output_to_dicts(
+    attributions: Any,
+    mask_type: str = "edge",
+    *args,
+    **kwargs,
+):
+    r"""Converts Captum attributions into dictionaries."""
+    raise NotImplementedError(
+        "Captum integration is PyTorch-specific and requires native torch.nn.Module."
+    )
+

k3_node/models/chemistry/__init__.py

@@ -0,0 +1,4 @@
+"""Chemistry and molecular graph neural network models (aliased from k3_node.applications.chemistry)."""
+
+from k3_node.applications.chemistry import *
+from k3_node.applications.chemistry import __all__

k3_node/models/correct_and_smooth.py

@@ -0,0 +1,146 @@
+
+import keras
+from keras import ops
+
+from k3_node.models.label_prop import LabelPropagation
+
+
+class CorrectAndSmooth(keras.layers.Layer):
+    r"""The correct and smooth (C&S) post-processing model from the
+    `"Combining Label Propagation And Simple Models Out-performs Graph Neural
+    Networks" <https://arxiv.org/abs/2010.13993>`_ paper.
+
+    Args:
+        num_correction_layers (int): The number of propagations :math:`L_1`.
+        correction_alpha (float): The :math:`\alpha_1` coefficient.
+        num_smoothing_layers (int): The number of propagations :math:`L_2`.
+        smoothing_alpha (float): The :math:`\alpha_2` coefficient.
+        autoscale (bool, optional): If set to :obj:`True`, will automatically
+            determine the scaling factor :math:`\gamma`. (default: :obj:`True`)
+        scale (float, optional): The scaling factor :math:`\gamma`, in case
+            :obj:`autoscale = False`. (default: :obj:`1.0`)
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.models import CorrectAndSmooth
+
+        y_soft = np.random.rand(6, 3).astype("float32")  # base model's class probabilities
+        y_true = np.array([1, 0, 0, 2, 1, 1])
+        train_mask = np.array([True, False, True, False, True, False])
+        edge_index = np.array([[0, 1, 1, 2, 4, 5], [1, 0, 2, 1, 5, 4]])
+
+        model = CorrectAndSmooth(num_correction_layers=2, correction_alpha=0.5,
+                                 num_smoothing_layers=2, smoothing_alpha=0.5)
+        y_soft = model.correct(y_soft, y_true[train_mask], train_mask, edge_index)  # propagate residual errors
+        y_soft = model.smooth(y_soft, y_true[train_mask], train_mask, edge_index)  # propagate predictions
+        print(tuple(y_soft.shape))  # (6, 3)
+        ```
+    """
+    def __init__(
+        self,
+        num_correction_layers: int,
+        correction_alpha: float,
+        num_smoothing_layers: int,
+        smoothing_alpha: float,
+        autoscale: bool = True,
+        scale: float = 1.0,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.autoscale = autoscale
+        self.scale = scale
+
+        self.prop1 = LabelPropagation(num_correction_layers, correction_alpha)
+        self.prop2 = LabelPropagation(num_smoothing_layers, smoothing_alpha)
+
+    def build(self, input_shape=None):
+        self.built = True
+
+    def call(self, y_soft, *args, **kwargs):
+        y_soft = self.correct(y_soft, *args, **kwargs)
+        return self.smooth(y_soft, *args, **kwargs)
+
+    def correct(self, y_soft, y_true, mask, edge_index, edge_weight=None):
+        # Plain NumPy inputs cannot be mixed with backend tensors; convert them first.
+        y_soft, y_true, mask = ops.convert_to_tensor(y_soft), ops.convert_to_tensor(y_true), ops.convert_to_tensor(mask)
+        num_classes = ops.shape(y_soft)[-1]
+        y_true_shape = ops.shape(y_true)
+        if len(y_true_shape) == 1:
+            y_true = ops.one_hot(y_true, num_classes)
+        y_true = ops.cast(y_true, y_soft.dtype)
+
+        mask_shape = ops.shape(mask)
+        is_bool_mask = (len(mask_shape) == 1 and 'bool' in str(mask.dtype))
+        if is_bool_mask:
+            indices = ops.where(mask)[0]
+        else:
+            indices = mask
+
+        numel = float(ops.shape(indices)[0])
+        if ops.shape(y_true)[0] == ops.shape(y_soft)[0]:
+            y_true_sub = ops.take(y_true, indices, axis=0)
+        else:
+            y_true_sub = y_true
+
+        error_zeros = ops.zeros_like(y_soft)
+        error = ops.scatter_update(
+            error_zeros,
+            ops.expand_dims(indices, -1),
+            y_true_sub - ops.take(y_soft, indices, axis=0)
+        )
+
+        if self.autoscale:
+            smoothed_error = self.prop1(
+                error, edge_index, edge_weight=edge_weight,
+                post_step=lambda x: ops.clip(x, -1.0, 1.0)
+            )
+
+            error_masked = ops.take(error, indices, axis=0)
+            sigma = ops.sum(ops.abs(error_masked)) / max(numel, 1.0)
+            sum_smoothed = ops.sum(ops.abs(smoothed_error), axis=1, keepdims=True)
+            scale = sigma / ops.maximum(sum_smoothed, 1e-12)
+            scale = ops.where(scale > 1000.0, ops.ones_like(scale), scale)
+            return y_soft + scale * smoothed_error
+        else:
+            def fix_input(x):
+                return ops.scatter_update(x, ops.expand_dims(indices, -1), ops.take(error, indices, axis=0))
+
+            smoothed_error = self.prop1(
+                error, edge_index, edge_weight=edge_weight,
+                post_step=fix_input,
+            )
+            return y_soft + self.scale * smoothed_error
+
+    def smooth(self, y_soft, y_true, mask, edge_index, edge_weight=None):
+        # Plain NumPy inputs cannot be mixed with backend tensors; convert them first.
+        y_soft, y_true, mask = ops.convert_to_tensor(y_soft), ops.convert_to_tensor(y_true), ops.convert_to_tensor(mask)
+        num_classes = ops.shape(y_soft)[-1]
+        y_true_shape = ops.shape(y_true)
+        if len(y_true_shape) == 1:
+            y_true = ops.one_hot(y_true, num_classes)
+        y_true = ops.cast(y_true, y_soft.dtype)
+
+        mask_shape = ops.shape(mask)
+        is_bool_mask = (len(mask_shape) == 1 and 'bool' in str(mask.dtype))
+        if is_bool_mask:
+            indices = ops.where(mask)[0]
+        else:
+            indices = mask
+
+        if ops.shape(y_true)[0] == ops.shape(y_soft)[0]:
+            y_true_sub = ops.take(y_true, indices, axis=0)
+        else:
+            y_true_sub = y_true
+
+        y_soft = ops.scatter_update(y_soft, ops.expand_dims(indices, -1), y_true_sub)
+        return self.prop2(y_soft, edge_index, edge_weight=edge_weight)
+
+    def __repr__(self):
+        L1, alpha1 = self.prop1.num_layers, self.prop1.alpha
+        L2, alpha2 = self.prop2.num_layers, self.prop2.alpha
+        return (f'{self.__class__.__name__}(\n'
+                f'  correct: num_layers={L1}, alpha={alpha1}\n'
+                f'  smooth:  num_layers={L2}, alpha={alpha2}\n'
+                f'  autoscale={self.autoscale}, scale={self.scale}\n'
+                ')')