k3-node 1.0.0__py3-none-any.whl

k3_node/layers/kge/transe.py

@@ -0,0 +1,96 @@
+import math
+
+import keras
+from keras import ops
+
+from k3_node.layers.kge.base import KGEModel, margin_ranking_loss, normalize
+
+
+class TransE(KGEModel):
+    r"""The TransE model from the `"Translating Embeddings for Modeling
+    Multi-Relational Data" <https://proceedings.neurips.cc/paper/2013/file/
+    1cecc7a77928ca8133fa24680a88d2f9-Paper.pdf>`_ paper.
+
+    :class:`TransE` models relations as a translation from head to tail
+    entities such that
+
+    .. math::
+        \mathbf{e}_h + \mathbf{e}_r \approx \mathbf{e}_t,
+
+    resulting in the scoring function:
+
+    .. math::
+        d(h, r, t) = - {\| \mathbf{e}_h + \mathbf{e}_r - \mathbf{e}_t \|}_p
+
+    Args:
+        num_nodes (int): The number of nodes/entities in the graph.
+        num_relations (int): The number of relations in the graph.
+        hidden_channels (int): The hidden embedding size.
+        margin (float, optional): The margin of the ranking loss.
+            (default: :obj:`1.0`)
+        p_norm (float, optional): The order embedding and distance
+            normalization. (default: :obj:`1.0`)
+        sparse (bool, optional): Kept for API compatibility. (default: :obj:`False`)
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import TransE
+
+        head = np.random.randint(0, 20, size=(10,))  # 10 (head, relation, tail) triples
+        rel = np.random.randint(0, 5, size=(10,))
+        tail = np.random.randint(0, 20, size=(10,))
+
+        model = TransE(num_nodes=20, num_relations=5, hidden_channels=8)
+        score = model(head, rel, tail)  # plausibility score of every triple
+        print(tuple(score.shape))  # (10,)
+        loss = model.loss(head, rel, tail)  # training loss against randomly corrupted triples
+        print(tuple(loss.shape))  # (): a scalar
+        ```
+    """
+    def __init__(
+        self,
+        num_nodes: int,
+        num_relations: int,
+        hidden_channels: int,
+        margin: float = 1.0,
+        p_norm: float = 1.0,
+        sparse: bool = False,
+        **kwargs,
+    ):
+        super().__init__(num_nodes, num_relations, hidden_channels, sparse, **kwargs)
+
+        self.p_norm = p_norm
+        self.margin = margin
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        bound = 6.0 / math.sqrt(self.hidden_channels)
+        # A new initializer per tensor: a reused unseeded Keras 3 initializer returns the same values on every call.
+        uniform = lambda shape: keras.initializers.RandomUniform(-bound, bound)(shape)
+        self.node_emb.embeddings.assign(uniform(ops.shape(self.node_emb.embeddings)))
+        self.rel_emb.embeddings.assign(uniform(ops.shape(self.rel_emb.embeddings)))
+        self.rel_emb.embeddings.assign(normalize(self.rel_emb.embeddings, p=self.p_norm, axis=-1))
+
+    def call(self, head_index, rel_type, tail_index):
+        head_index = ops.cast(head_index, "int32")
+        rel_type = ops.cast(rel_type, "int32")
+        tail_index = ops.cast(tail_index, "int32")
+
+        head = self.node_emb(head_index)
+        rel = self.rel_emb(rel_type)
+        tail = self.node_emb(tail_index)
+
+        head = normalize(head, p=self.p_norm, axis=-1)
+        tail = normalize(tail, p=self.p_norm, axis=-1)
+
+        # Calculate *negative* TransE norm:
+        diff = (head + rel) - tail
+        return -ops.power(ops.sum(ops.power(ops.abs(diff), self.p_norm), axis=-1), 1.0 / self.p_norm)
+
+    def loss(self, head_index, rel_type, tail_index):
+        pos_score = self(head_index, rel_type, tail_index)
+        neg_score = self(*self.random_sample(head_index, rel_type, tail_index))
+
+        return margin_ranking_loss(pos_score, neg_score, margin=self.margin)

k3_node/layers/norm/__init__.py

@@ -0,0 +1,23 @@
+from .batch_norm import BatchNorm, HeteroBatchNorm
+from .diff_group_norm import DiffGroupNorm
+from .graph_norm import GraphNorm
+from .graph_size_norm import GraphSizeNorm
+from .instance_norm import InstanceNorm
+from .layer_norm import HeteroLayerNorm, LayerNorm
+from .mean_subtraction_norm import MeanSubtractionNorm
+from .msg_norm import MessageNorm
+from .pair_norm import PairNorm
+
+__all__ = [
+    "BatchNorm",
+    "HeteroBatchNorm",
+    "InstanceNorm",
+    "LayerNorm",
+    "HeteroLayerNorm",
+    "GraphNorm",
+    "GraphSizeNorm",
+    "PairNorm",
+    "MeanSubtractionNorm",
+    "MessageNorm",
+    "DiffGroupNorm",
+]

k3_node/layers/norm/batch_norm.py

@@ -0,0 +1,328 @@
+from typing import Optional
+from keras import layers, ops
+from k3_node.ops.segment import segment_sum
+
+
+class BatchNorm(layers.Layer):
+    r"""Applies batch normalization over a batch of features as described in
+    the `"Batch Normalization: Accelerating Deep Network Training by
+    Reducing Internal Covariate Shift" <https://arxiv.org/abs/1502.03167>`_
+    paper.
+
+    .. math::
+        \mathbf{x}^{\prime}_i = \frac{\mathbf{x} -
+        \textrm{E}[\mathbf{x}]}{\sqrt{\textrm{Var}[\mathbf{x}] + \epsilon}}
+        \odot \gamma + \beta
+
+    Args:
+        in_channels (int): Size of each input sample.
+        eps (float, optional): A value added to the denominator for numerical
+            stability. (default: :obj:`1e-5`)
+        momentum (float, optional): The value used for the running mean and
+            running variance computation. (default: :obj:`0.1`)
+        affine (bool, optional): If set to :obj:`True`, this module has
+            learnable affine parameters :math:`\gamma` and :math:`\beta`.
+            (default: :obj:`True`)
+        track_running_stats (bool, optional): If set to :obj:`True`, this
+            module tracks the running mean and variance, and when set to
+            :obj:`False`, this module does not track such statistics and always
+            uses batch statistics in both training and eval modes.
+            (default: :obj:`True`)
+        allow_single_element (bool, optional): If set to :obj:`True`, batches
+            with only a single element will work as during in evaluation.
+            That is the running mean and variance will be used.
+            Requires :obj:`track_running_stats=True`. (default: :obj:`False`)
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import BatchNorm
+
+        x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
+
+        layer = BatchNorm(in_channels=8)
+        out = layer(x, training=True)  # uses batch statistics while training
+        print(tuple(out.shape))  # (10, 8)
+        ```
+    """
+    def __init__(
+        self,
+        in_channels: int,
+        eps: float = 1e-5,
+        momentum: Optional[float] = 0.1,
+        affine: bool = True,
+        track_running_stats: bool = True,
+        allow_single_element: bool = False,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        if allow_single_element and not track_running_stats:
+            raise ValueError("'allow_single_element' requires "
+                             "'track_running_stats' to be set to `True`")
+
+        self.in_channels = in_channels
+        self.eps = eps
+        self.momentum = momentum
+        self.affine = affine
+        self.track_running_stats = track_running_stats
+        self.allow_single_element = allow_single_element
+
+        if affine:
+            self.weight = self.add_weight(
+                shape=(in_channels,),
+                initializer="ones",
+                trainable=True,
+                name="weight",
+            )
+            self.bias = self.add_weight(
+                shape=(in_channels,),
+                initializer="zeros",
+                trainable=True,
+                name="bias",
+            )
+        else:
+            self.weight = None
+            self.bias = None
+
+        if track_running_stats:
+            self.running_mean = self.add_weight(
+                shape=(in_channels,),
+                initializer="zeros",
+                trainable=False,
+                name="running_mean",
+            )
+            self.running_var = self.add_weight(
+                shape=(in_channels,),
+                initializer="ones",
+                trainable=False,
+                name="running_var",
+            )
+            self.num_batches_tracked = self.add_weight(
+                shape=(),
+                initializer="zeros",
+                dtype="int32",
+                trainable=False,
+                name="num_batches_tracked",
+            )
+        else:
+            self.running_mean = None
+            self.running_var = None
+            self.num_batches_tracked = None
+
+    def reset_running_stats(self):
+        if self.track_running_stats:
+            self.running_mean.assign(ops.zeros(self.running_mean.shape, dtype=self.running_mean.dtype))
+            self.running_var.assign(ops.ones(self.running_var.shape, dtype=self.running_var.dtype))
+            self.num_batches_tracked.assign(ops.cast(0, "int32"))
+
+    def reset_parameters(self):
+        self.reset_running_stats()
+        if self.affine:
+            self.weight.assign(ops.ones(self.weight.shape, dtype=self.weight.dtype))
+            self.bias.assign(ops.zeros(self.bias.shape, dtype=self.bias.dtype))
+
+    def call(self, x, training=None):
+        num_samples_static = x.shape[0]
+        # Keras semantics: `training=None` means inference; fit() passes training=True via the call context.
+        is_training = bool(training) if training is not None else False
+
+        if is_training:
+            if num_samples_static is not None and num_samples_static <= 1:
+                if not self.allow_single_element:
+                    raise ValueError(f"Expected more than 1 value per channel when training, got input size {ops.shape(x)}")
+                # Evaluation behavior with running stats
+                mean = self.running_mean
+                var = self.running_var
+            else:
+                mean = ops.mean(x, axis=0)
+                var = ops.var(x, axis=0)
+
+                if self.track_running_stats:
+                    n = ops.cast(ops.shape(x)[0], dtype=x.dtype)
+                    unbiased_var = var * n / ops.maximum(n - 1.0, 1.0)
+                    if self.momentum is None:
+                        count = ops.cast(self.num_batches_tracked + 1, dtype=x.dtype)
+                        m = 1.0 / count
+                    else:
+                        m = self.momentum
+
+                    new_running_mean = (1.0 - m) * self.running_mean + m * mean
+                    new_running_var = (1.0 - m) * self.running_var + m * unbiased_var
+                    self.running_mean.assign(new_running_mean)
+                    self.running_var.assign(new_running_var)
+                    self.num_batches_tracked.assign(self.num_batches_tracked + 1)
+        else:
+            if self.track_running_stats:
+                mean = self.running_mean
+                var = self.running_var
+            else:
+                mean = ops.mean(x, axis=0)
+                var = ops.var(x, axis=0)
+
+        out = (x - mean) / ops.sqrt(var + self.eps)
+
+        if self.affine:
+            out = out * self.weight + self.bias
+
+        return out
+
+    def compute_output_shape(self, input_shape):
+        return input_shape
+
+
+class HeteroBatchNorm(layers.Layer):
+    r"""Applies batch normalization over a batch of heterogeneous features as
+    described in the `"Batch Normalization: Accelerating Deep Network Training
+    by Reducing Internal Covariate Shift" <https://arxiv.org/abs/1502.03167>`_
+    paper.
+    Compared to :class:`BatchNorm`, :class:`HeteroBatchNorm` applies
+    normalization individually for each node or edge type.
+
+    Args:
+        in_channels (int): Size of each input sample.
+        num_types (int): The number of types.
+        eps (float, optional): A value added to the denominator for numerical
+            stability. (default: :obj:`1e-5`)
+        momentum (float, optional): The value used for the running mean and
+            running variance computation. (default: :obj:`0.1`)
+        affine (bool, optional): If set to :obj:`True`, this module has
+            learnable affine parameters :math:`\gamma` and :math:`\beta`.
+            (default: :obj:`True`)
+        track_running_stats (bool, optional): If set to :obj:`True`, this
+            module tracks the running mean and variance, and when set to
+            :obj:`False`, this module does not track such statistics and always
+            uses batch statistics in both training and eval modes.
+            (default: :obj:`True`)
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import HeteroBatchNorm
+
+        x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
+        node_type = np.random.randint(0, 3, size=(10,))  # type of each node
+
+        layer = HeteroBatchNorm(in_channels=8, num_types=3)
+        out = layer(x, node_type, training=True)  # separate statistics per node type
+        print(tuple(out.shape))  # (10, 8)
+        ```
+    """
+    def __init__(
+        self,
+        in_channels: int,
+        num_types: int,
+        eps: float = 1e-5,
+        momentum: Optional[float] = 0.1,
+        affine: bool = True,
+        track_running_stats: bool = True,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.in_channels = in_channels
+        self.num_types = num_types
+        self.eps = eps
+        self.momentum = momentum
+        self.affine = affine
+        self.track_running_stats = track_running_stats
+
+        if affine:
+            self.weight = self.add_weight(
+                shape=(num_types, in_channels),
+                initializer="ones",
+                trainable=True,
+                name="weight",
+            )
+            self.bias = self.add_weight(
+                shape=(num_types, in_channels),
+                initializer="zeros",
+                trainable=True,
+                name="bias",
+            )
+        else:
+            self.weight = None
+            self.bias = None
+
+        if track_running_stats:
+            self.running_mean = self.add_weight(
+                shape=(num_types, in_channels),
+                initializer="zeros",
+                trainable=False,
+                name="running_mean",
+            )
+            self.running_var = self.add_weight(
+                shape=(num_types, in_channels),
+                initializer="ones",
+                trainable=False,
+                name="running_var",
+            )
+            self.num_batches_tracked = self.add_weight(
+                shape=(),
+                initializer="zeros",
+                dtype="int32",
+                trainable=False,
+                name="num_batches_tracked",
+            )
+        else:
+            self.running_mean = None
+            self.running_var = None
+            self.num_batches_tracked = None
+
+    def reset_running_stats(self):
+        if self.track_running_stats:
+            self.running_mean.assign(ops.zeros(self.running_mean.shape, dtype=self.running_mean.dtype))
+            self.running_var.assign(ops.ones(self.running_var.shape, dtype=self.running_var.dtype))
+            self.num_batches_tracked.assign(ops.cast(0, "int32"))
+
+    def reset_parameters(self):
+        self.reset_running_stats()
+        if self.affine:
+            self.weight.assign(ops.ones(self.weight.shape, dtype=self.weight.dtype))
+            self.bias.assign(ops.zeros(self.bias.shape, dtype=self.bias.dtype))
+
+    def call(self, x, type_vec=None, training=None):
+        if type_vec is None and isinstance(x, (tuple, list)):
+            x, type_vec = x
+
+        # Keras semantics: `training=None` means inference; fit() passes training=True via the call context.
+        is_training = bool(training) if training is not None else False
+        type_vec = ops.cast(type_vec, "int32")
+
+        if not is_training and self.track_running_stats:
+            mean = self.running_mean
+            var = self.running_var
+        else:
+            ones = ops.ones((ops.shape(x)[0], 1), dtype=x.dtype)
+            counts = ops.maximum(segment_sum(ones, type_vec, num_segments=self.num_types), 1.0)
+            mean = segment_sum(x, type_vec, num_segments=self.num_types) / counts
+            x_c = x - ops.take(mean, type_vec, axis=0)
+            var = segment_sum(ops.power(x_c, 2), type_vec, num_segments=self.num_types) / counts
+
+        if is_training and self.track_running_stats:
+            if self.momentum is None:
+                count = ops.cast(self.num_batches_tracked + 1, dtype=x.dtype)
+                exp_avg_factor = 1.0 / count
+            else:
+                exp_avg_factor = self.momentum
+
+            new_running_mean = (1.0 - exp_avg_factor) * self.running_mean + exp_avg_factor * mean
+            new_running_var = (1.0 - exp_avg_factor) * self.running_var + exp_avg_factor * var
+            self.running_mean.assign(new_running_mean)
+            self.running_var.assign(new_running_var)
+            self.num_batches_tracked.assign(self.num_batches_tracked + 1)
+
+        std = ops.sqrt(var + self.eps)
+        mean_taken = ops.take(mean, type_vec, axis=0)
+        std_taken = ops.take(std, type_vec, axis=0)
+        out = (x - mean_taken) / std_taken
+
+        if self.affine:
+            w_taken = ops.take(self.weight, type_vec, axis=0)
+            b_taken = ops.take(self.bias, type_vec, axis=0)
+            out = out * w_taken + b_taken
+
+        return out
+
+    def compute_output_shape(self, input_shape):
+        if isinstance(input_shape, (tuple, list)):
+            return input_shape[0]
+        return input_shape

k3_node/layers/norm/diff_group_norm.py

@@ -0,0 +1,141 @@
+import numpy as np
+from scipy.spatial.distance import cdist
+from keras import initializers, layers, ops
+
+from .batch_norm import BatchNorm
+
+
+class DiffGroupNorm(layers.Layer):
+    r"""The differentiable group normalization layer from the `"Towards Deeper
+    Graph Neural Networks with Differentiable Group Normalization"
+    <https://arxiv.org/abs/2006.06972>`_ paper, which normalizes node features
+    group-wise via a learnable soft cluster assignment.
+
+    .. math::
+
+        \mathbf{S} = \text{softmax} (\mathbf{X} \mathbf{W})
+
+    where :math:`\mathbf{W} \in \mathbb{R}^{F \times G}` denotes a trainable
+    weight matrix mapping each node into one of :math:`G` clusters.
+    Normalization is then performed group-wise via:
+
+    .. math::
+
+        \mathbf{X}^{\prime} = \mathbf{X} + \lambda \sum_{i = 1}^G
+        \text{BatchNorm}(\mathbf{S}[:, i] \odot \mathbf{X})
+
+    Args:
+        in_channels (int): Size of each input sample :math:`F`.
+        groups (int): The number of groups :math:`G`.
+        lamda (float, optional): The balancing factor :math:`\lambda` between
+            input embeddings and normalized embeddings. (default: :obj:`0.01`)
+        eps (float, optional): A value added to the denominator for numerical
+            stability. (default: :obj:`1e-5`)
+        momentum (float, optional): The value used for the running mean and
+            running variance computation. (default: :obj:`0.1`)
+        affine (bool, optional): If set to :obj:`True`, this module has
+            learnable affine parameters :math:`\gamma` and :math:`\beta`.
+            (default: :obj:`True`)
+        track_running_stats (bool, optional): If set to :obj:`True`, this
+            module tracks the running mean and variance, and when set to
+            :obj:`False`, this module does not track such statistics and always
+            uses batch statistics in both training and eval modes.
+            (default: :obj:`True`)
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import DiffGroupNorm
+
+        x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
+
+        layer = DiffGroupNorm(in_channels=8, groups=2)
+        out = layer(x, training=True)  # uses batch statistics while training
+        print(tuple(out.shape))  # (10, 8)
+        ```
+    """
+    def __init__(
+        self,
+        in_channels: int,
+        groups: int,
+        lamda: float = 0.01,
+        eps: float = 1e-5,
+        momentum: float = 0.1,
+        affine: bool = True,
+        track_running_stats: bool = True,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.in_channels = in_channels
+        self.groups = groups
+        self.lamda = lamda
+        self.eps = eps
+        self.momentum = momentum
+        self.affine = affine
+        self.track_running_stats = track_running_stats
+
+        self.lin_weight = self.add_weight(
+            shape=(in_channels, groups),
+            initializer="glorot_uniform",
+            trainable=True,
+            name="lin_weight",
+        )
+        self.norm = BatchNorm(
+            groups * in_channels,
+            eps=eps,
+            momentum=momentum,
+            affine=affine,
+            track_running_stats=track_running_stats,
+            name="norm",
+        )
+
+    def reset_parameters(self):
+        self.lin_weight.assign(
+            initializers.GlorotUniform()(shape=self.lin_weight.shape, dtype=self.lin_weight.dtype)
+        )
+        self.norm.reset_parameters()
+
+    def call(self, x, training=None):
+        F, G = self.in_channels, self.groups
+
+        s = ops.softmax(ops.matmul(x, self.lin_weight), axis=-1)  # [N, G]
+        out = ops.expand_dims(s, axis=-1) * ops.expand_dims(x, axis=-2)  # [N, G, F]
+        out_flat = ops.reshape(out, (-1, G * F))
+        out_norm = self.norm(out_flat, training=training)
+        out = ops.sum(ops.reshape(out_norm, (-1, G, F)), axis=-2)  # [N, F]
+
+        return x + self.lamda * out
+
+    @staticmethod
+    def group_distance_ratio(x, y, eps: float = 1e-5) -> float:
+        r"""Measures the ratio of inter-group distance over intra-group
+        distance.
+        """
+        x_np = ops.convert_to_numpy(x)
+        y_np = ops.convert_to_numpy(y).astype(np.int64)
+
+        num_classes = int(y_np.max()) + 1
+
+        numerator = 0.0
+        for i in range(num_classes):
+            mask = (y_np == i)
+            if not np.any(mask) or np.all(mask):
+                continue
+            dist = cdist(x_np[mask], x_np[~mask])
+            numerator += (1.0 / dist.size) * float(dist.sum())
+        numerator *= 1.0 / ((num_classes - 1) ** 2)
+
+        denominator = 0.0
+        for i in range(num_classes):
+            mask = (y_np == i)
+            if not np.any(mask):
+                continue
+            dist = cdist(x_np[mask], x_np[mask])
+            denominator += (1.0 / dist.size) * float(dist.sum())
+        denominator *= 1.0 / num_classes
+
+        return float(numerator / (denominator + eps))
+
+    def compute_output_shape(self, input_shape):
+        return input_shape
+

k3_node/layers/norm/graph_norm.py

@@ -0,0 +1,105 @@
+from keras import layers, ops
+from k3_node.layers.conv.utils import is_tracing
+from k3_node.ops.segment import segment_sum
+
+
+class GraphNorm(layers.Layer):
+    r"""Applies graph normalization over individual graphs as described in the
+    `"GraphNorm: A Principled Approach to Accelerating Graph Neural Network
+    Training" <https://arxiv.org/abs/2009.03294>`_ paper.
+
+    .. math::
+        \mathbf{x}^{\prime}_i = \frac{\mathbf{x} - \alpha \odot
+        \textrm{E}[\mathbf{x}]}
+        {\sqrt{\textrm{Var}[\mathbf{x} - \alpha \odot \textrm{E}[\mathbf{x}]]
+        + \epsilon}} \odot \gamma + \beta
+
+    where :math:`\alpha` denotes parameters that learn how much information
+    to keep in the mean.
+
+    Args:
+        in_channels (int): Size of each input sample.
+        eps (float, optional): A value added to the denominator for numerical
+            stability. (default: :obj:`1e-5`)
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import GraphNorm
+
+        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
+
+        layer = GraphNorm(in_channels=8)
+        out = layer(x, batch)  # normalizes each graph separately
+        print(tuple(out.shape))  # (10, 8)
+        ```
+    """
+    def __init__(self, in_channels: int, eps: float = 1e-5, **kwargs):
+        super().__init__(**kwargs)
+        self.in_channels = in_channels
+        self.eps = eps
+
+        self.weight = self.add_weight(
+            shape=(in_channels,),
+            initializer="ones",
+            trainable=True,
+            name="weight",
+        )
+        self.bias = self.add_weight(
+            shape=(in_channels,),
+            initializer="zeros",
+            trainable=True,
+            name="bias",
+        )
+        self.mean_scale = self.add_weight(
+            shape=(in_channels,),
+            initializer="ones",
+            trainable=True,
+            name="mean_scale",
+        )
+
+    def reset_parameters(self):
+        self.weight.assign(ops.ones(self.weight.shape, dtype=self.weight.dtype))
+        self.bias.assign(ops.zeros(self.bias.shape, dtype=self.bias.dtype))
+        self.mean_scale.assign(ops.ones(self.mean_scale.shape, dtype=self.mean_scale.dtype))
+
+    def call(self, x, batch=None, batch_size=None):
+        if batch is None and isinstance(x, (tuple, list)):
+            if len(x) == 2:
+                x, batch = x
+            elif len(x) == 3:
+                x, batch, batch_size = x
+
+        if batch is None:
+            mean = ops.mean(x, axis=0, keepdims=True)
+            out = x - mean * self.mean_scale
+            var = ops.mean(ops.power(out, 2), axis=0, keepdims=True)
+            std = ops.sqrt(var + self.eps)
+            return self.weight * out / std + self.bias
+
+        if batch_size is None:
+            if not is_tracing(batch):
+                try:
+                    batch_size = int(ops.max(batch)) + 1
+                except Exception:
+                    batch_size = None
+        elif not isinstance(batch_size, int):
+            try:
+                batch_size = int(batch_size)
+            except Exception:
+                pass
+
+        batch = ops.cast(batch, "int32")
+        ones = ops.ones((ops.shape(x)[0], 1), dtype=x.dtype)
+        counts = ops.maximum(segment_sum(ones, batch, num_segments=batch_size), 1.0)
+        mean = segment_sum(x, batch, num_segments=batch_size) / counts
+        out = x - ops.take(mean, batch, axis=0) * self.mean_scale
+        var = segment_sum(ops.power(out, 2), batch, num_segments=batch_size) / counts
+        std = ops.take(ops.sqrt(var + self.eps), batch, axis=0)
+        return self.weight * out / std + self.bias
+
+    def compute_output_shape(self, input_shape):
+        if isinstance(input_shape, (tuple, list)) and isinstance(input_shape[0], (tuple, list)):
+            return input_shape[0]
+        return input_shape