k3-node 1.0.0__py3-none-any.whl

k3_node/layers/norm/graph_size_norm.py

@@ -0,0 +1,57 @@
+from keras import layers, ops
+from k3_node.ops.segment import segment_sum
+
+
+class GraphSizeNorm(layers.Layer):
+    r"""Applies Graph Size Normalization over each individual graph in a batch
+    of node features:
+
+    .. math::
+        \mathbf{x}^{\prime}_i = \frac{\mathbf{x}_i}{\sqrt{|\mathcal{V}|}}
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import GraphSizeNorm
+
+        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 = GraphSizeNorm()
+        out = layer(x, batch)  # normalizes each graph separately
+        print(tuple(out.shape))  # (10, 8)
+        ```
+    """
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+    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:
+            num_nodes = ops.cast(ops.shape(x)[0], dtype=x.dtype)
+            return x * ops.power(num_nodes, -0.5)
+
+        if batch_size is not None and not isinstance(batch_size, int):
+            try:
+                batch_size = int(batch_size)
+            except Exception:
+                pass
+        elif batch_size is None:
+            batch_size = ops.cast(ops.max(batch), "int32") + 1
+
+        batch = ops.cast(batch, "int32")
+        ones = ops.ones((ops.shape(x)[0], 1), dtype=x.dtype)
+        deg = segment_sum(ones, batch, num_segments=batch_size)
+        inv_sqrt_deg = ops.power(deg, -0.5)
+        scale = ops.take(inv_sqrt_deg, batch, axis=0)
+        return x * scale
+
+    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

k3_node/layers/norm/instance_norm.py

@@ -0,0 +1,163 @@
+from keras import layers, ops
+from k3_node.ops.segment import segment_sum
+
+
+class InstanceNorm(layers.Layer):
+    r"""Applies instance normalization over each individual example in a batch
+    of node features as described in the `"Instance Normalization: The Missing
+    Ingredient for Fast Stylization" <https://arxiv.org/abs/1607.06450>`_
+    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:`False`)
+        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 instance statistics in both training and eval modes.
+            (default: :obj:`False`)
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import InstanceNorm
+
+        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 = InstanceNorm(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,
+        momentum: float = 0.1,
+        affine: bool = False,
+        track_running_stats: bool = False,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.in_channels = in_channels
+        self.eps = eps
+        self.momentum = momentum
+        self.affine = affine
+        self.track_running_stats = track_running_stats
+
+        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",
+            )
+        else:
+            self.running_mean = None
+            self.running_var = 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))
+
+    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, batch=None, batch_size=None, training=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
+
+        # 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 batch is None:
+            batch = ops.zeros((ops.shape(x)[0],), dtype="int32")
+            batch_size = 1
+        elif batch_size is not None and not isinstance(batch_size, int):
+            try:
+                batch_size = int(batch_size)
+            except Exception:
+                pass
+        elif batch_size is None:
+            batch_size = ops.cast(ops.max(batch), "int32") + 1
+
+        batch = ops.cast(batch, "int32")
+
+        if is_training or not self.track_running_stats:
+            ones = ops.ones((ops.shape(x)[0], 1), dtype=x.dtype)
+            counts = ops.maximum(segment_sum(ones, batch, num_segments=batch_size), 1.0)
+            unbiased_counts = ops.maximum(counts - 1.0, 1.0)
+
+            mean = segment_sum(x, batch, num_segments=batch_size) / counts
+            x_c = x - ops.take(mean, batch, axis=0)
+            sq_diff = segment_sum(ops.power(x_c, 2), batch, num_segments=batch_size)
+            var = sq_diff / counts
+            unbiased_var = sq_diff / unbiased_counts
+
+            if is_training and self.track_running_stats:
+                m = self.momentum
+                cur_mean = ops.mean(mean, axis=0)
+                cur_var = ops.mean(unbiased_var, axis=0)
+                new_running_mean = (1.0 - m) * self.running_mean + m * cur_mean
+                new_running_var = (1.0 - m) * self.running_var + m * cur_var
+                self.running_mean.assign(new_running_mean)
+                self.running_var.assign(new_running_var)
+
+            std_x = ops.take(ops.sqrt(var + self.eps), batch, axis=0)
+            out = x_c / std_x
+        else:
+            x_c = x - self.running_mean
+            std = ops.sqrt(self.running_var + self.eps)
+            out = x_c / std
+
+        if self.affine:
+            out = out * self.weight + self.bias
+
+        return out
+
+    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

k3_node/layers/norm/layer_norm.py

@@ -0,0 +1,245 @@
+from typing import List, Optional, Union
+from keras import layers, ops
+from k3_node.ops.segment import segment_sum
+
+
+class LayerNorm(layers.Layer):
+    r"""Applies layer normalization over each individual example in a batch
+    of features as described in the `"Layer Normalization"
+    <https://arxiv.org/abs/1607.06450>`_ 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`)
+        affine (bool, optional): If set to :obj:`True`, this module has
+            learnable affine parameters :math:`\gamma` and :math:`\beta`.
+            (default: :obj:`True`)
+        mode (str, optional): The normalization mode to use for layer
+            normalization (:obj:`"graph"` or :obj:`"node"`). If :obj:`"graph"`
+            is used, each graph will be considered as an element to be
+            normalized. If `"node"` is used, each node will be considered as
+            an element to be normalized. (default: :obj:`"graph"`)
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import LayerNorm
+
+        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 = LayerNorm(in_channels=8, mode="graph")
+        out = layer(x, batch)  # normalizes each graph separately
+        print(tuple(out.shape))  # (10, 8)
+        ```
+    """
+    def __init__(
+        self,
+        in_channels: int,
+        eps: float = 1e-5,
+        affine: bool = True,
+        mode: str = 'graph',
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        if mode not in ('graph', 'node'):
+            raise ValueError(f"Unknown normalization mode: {mode}")
+
+        self.in_channels = in_channels
+        self.eps = eps
+        self.affine = affine
+        self.mode = mode
+
+        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
+
+    def reset_parameters(self):
+        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, 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 self.mode == 'graph':
+            if batch is None:
+                mean = ops.mean(x)
+                var = ops.mean(ops.power(x - mean, 2))
+                out = (x - mean) / ops.sqrt(var + self.eps)
+            else:
+                if batch_size is not None and not isinstance(batch_size, int):
+                    try:
+                        batch_size = int(batch_size)
+                    except Exception:
+                        pass
+                elif batch_size is None:
+                    batch_size = ops.cast(ops.max(batch), "int32") + 1
+
+                batch = ops.cast(batch, "int32")
+                in_channels = ops.cast(ops.shape(x)[-1], dtype=x.dtype)
+                ones = ops.ones((ops.shape(x)[0], 1), dtype=x.dtype)
+                node_counts = ops.maximum(segment_sum(ones, batch, num_segments=batch_size), 1.0)
+                total_count = node_counts * in_channels
+
+                sum_x = ops.sum(segment_sum(x, batch, num_segments=batch_size), axis=-1, keepdims=True)
+                mean = sum_x / total_count
+                x_centered = x - ops.take(mean, batch, axis=0)
+
+                sum_sq = ops.sum(segment_sum(ops.power(x_centered, 2), batch, num_segments=batch_size), axis=-1, keepdims=True)
+                var = sum_sq / total_count
+                std_x = ops.take(ops.sqrt(var + self.eps), batch, axis=0)
+                out = x_centered / std_x
+
+            if self.affine:
+                out = out * self.weight + self.bias
+            return out
+
+        elif self.mode == 'node':
+            mean = ops.mean(x, axis=-1, keepdims=True)
+            var = ops.var(x, axis=-1, keepdims=True)
+            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):
+        if isinstance(input_shape, (tuple, list)) and isinstance(input_shape[0], (tuple, list)):
+            return input_shape[0]
+        return input_shape
+
+
+class HeteroLayerNorm(layers.Layer):
+    r"""Applies layer normalization over each individual example in a batch
+    of heterogeneous features as described in the `"Layer Normalization"
+    <https://arxiv.org/abs/1607.06450>`_ paper.
+    Compared to :class:`LayerNorm`, :class:`HeteroLayerNorm` 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`)
+        affine (bool, optional): If set to :obj:`True`, this module has
+            learnable affine parameters :math:`\gamma` and :math:`\beta`.
+            (default: :obj:`True`)
+        mode (str, optional): The normalization mode to use for layer
+            normalization (:obj:`"node"`). (default: :obj:`"node"`)
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import HeteroLayerNorm
+
+        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 = HeteroLayerNorm(in_channels=8, num_types=3)
+        out = layer(x, node_type)  # separate statistics per node type
+        print(tuple(out.shape))  # (10, 8)
+        ```
+    """
+    def __init__(
+        self,
+        in_channels: int,
+        num_types: int,
+        eps: float = 1e-5,
+        affine: bool = True,
+        mode: str = 'node',
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        if mode != 'node':
+            raise ValueError(f"HeteroLayerNorm only supports mode='node' (got '{mode}')")
+
+        self.in_channels = in_channels
+        self.num_types = num_types
+        self.eps = eps
+        self.affine = affine
+        self.mode = mode
+
+        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
+
+    def reset_parameters(self):
+        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,
+        type_ptr: Optional[Union[list, tuple]] = None,
+    ):
+        if type_vec is None and isinstance(x, (tuple, list)):
+            if len(x) == 2:
+                x, type_vec = x
+            elif len(x) == 3:
+                x, type_vec, type_ptr = x
+
+        if type_vec is None and type_ptr is None:
+            raise ValueError("Either 'type_vec' or 'type_ptr' must be given")
+
+        mean = ops.mean(x, axis=-1, keepdims=True)
+        var = ops.var(x, axis=-1, keepdims=True)
+        out = (x - mean) / ops.sqrt(var + self.eps)
+
+        if self.affine:
+            if type_ptr is not None:
+                parts = []
+                for i in range(len(type_ptr) - 1):
+                    s, e = type_ptr[i], type_ptr[i + 1]
+                    part = out[s:e] * self.weight[i] + self.bias[i]
+                    parts.append(part)
+                out = ops.concatenate(parts, axis=0)
+            else:
+                type_vec = ops.cast(type_vec, "int32")
+                w = ops.take(self.weight, type_vec, axis=0)
+                b = ops.take(self.bias, type_vec, axis=0)
+                out = out * w + b
+
+        return out
+
+    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

k3_node/layers/norm/mean_subtraction_norm.py

@@ -0,0 +1,57 @@
+from keras import layers, ops
+from k3_node.ops.segment import segment_sum
+
+
+class MeanSubtractionNorm(layers.Layer):
+    r"""Applies layer normalization by subtracting the mean from the inputs
+    as described in the `"Revisiting 'Over-smoothing' in Deep GCNs"
+    <https://arxiv.org/abs/2003.13663>`_ paper.
+
+    .. math::
+        \mathbf{x}_i = \mathbf{x}_i - \frac{1}{|\mathcal{V}|}
+        \sum_{j \in \mathcal{V}} \mathbf{x}_j
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import MeanSubtractionNorm
+
+        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 = MeanSubtractionNorm()
+        out = layer(x, batch)  # normalizes each graph separately
+        print(tuple(out.shape))  # (10, 8)
+        ```
+    """
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+    def call(self, x, batch=None, dim_size=None):
+        if batch is None and isinstance(x, (tuple, list)):
+            if len(x) == 2:
+                x, batch = x
+            elif len(x) == 3:
+                x, batch, dim_size = x
+
+        if batch is None:
+            return x - ops.mean(x, axis=0, keepdims=True)
+
+        if dim_size is not None and not isinstance(dim_size, int):
+            try:
+                dim_size = int(dim_size)
+            except Exception:
+                pass
+        elif dim_size is None:
+            dim_size = ops.cast(ops.max(batch), "int32") + 1
+
+        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=dim_size), 1.0)
+        mean = segment_sum(x, batch, num_segments=dim_size) / counts
+        return x - ops.take(mean, batch, axis=0)
+
+    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

k3_node/layers/norm/msg_norm.py

@@ -0,0 +1,58 @@
+from keras import layers, ops
+
+
+class MessageNorm(layers.Layer):
+    r"""Applies message normalization over the aggregated messages as described
+    in the `"DeeperGCNs: All You Need to Train Deeper GCNs"
+    <https://arxiv.org/abs/2006.07739>`_ paper.
+
+    .. math::
+
+        \mathbf{x}_i^{\prime} = \mathbf{x}_{i} + s \cdot
+        {\| \mathbf{x}_i \|}_2 \cdot
+        \frac{\mathbf{m}_{i}}{{\|\mathbf{m}_i\|}_2}
+
+    Args:
+        learn_scale (bool, optional): If set to :obj:`True`, will learn the
+            scaling factor :math:`s` of message normalization.
+            (default: :obj:`False`)
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import MessageNorm
+
+        x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
+
+        msg = np.random.rand(10, 8).astype("float32")  # aggregated messages for each node
+        layer = MessageNorm(learn_scale=True)
+        out = layer(x, msg)  # rescales messages to the norm of x
+        print(tuple(out.shape))  # (10, 8)
+        ```
+    """
+    def __init__(self, learn_scale: bool = False, **kwargs):
+        super().__init__(**kwargs)
+        self.learn_scale = learn_scale
+        self.scale = self.add_weight(
+            shape=(1,),
+            initializer="ones",
+            trainable=learn_scale,
+            name="scale",
+        )
+
+    def reset_parameters(self):
+        self.scale.assign(ops.ones(self.scale.shape, dtype=self.scale.dtype))
+
+    def call(self, x, msg=None, p=2.0):
+        if msg is None and isinstance(x, (tuple, list)):
+            x, msg = x
+
+        msg_norm = ops.maximum(ops.norm(msg, ord=p, axis=-1, keepdims=True), 1e-12)
+        msg_normalized = msg / msg_norm
+        x_norm = ops.norm(x, ord=p, axis=-1, keepdims=True)
+        return msg_normalized * x_norm * self.scale
+
+    def compute_output_shape(self, input_shape):
+        if isinstance(input_shape, (tuple, list)):
+            return input_shape[0]
+        return input_shape

k3_node/layers/norm/pair_norm.py

@@ -0,0 +1,94 @@
+from keras import layers, ops
+from k3_node.ops.segment import segment_sum
+
+
+class PairNorm(layers.Layer):
+    r"""Applies pair normalization over node features as described in the
+    `"PairNorm: Tackling Oversmoothing in GNNs"
+    <https://arxiv.org/abs/1909.12223>`_ paper.
+
+    .. math::
+        \mathbf{x}_i^c &= \mathbf{x}_i - \frac{1}{n}
+        \sum_{i=1}^n \mathbf{x}_i \\
+
+        \mathbf{x}_i^{\prime} &= s \cdot
+        \frac{\mathbf{x}_i^c}{\sqrt{\frac{1}{n} \sum_{i=1}^n
+        {\| \mathbf{x}_i^c \|}^2_2}}
+
+    Args:
+        scale (float, optional): Scaling factor :math:`s` of normalization.
+            (default: :obj:`1.0`)
+        scale_individually (bool, optional): If set to :obj:`True`, will
+            compute the scaling step as :math:`\mathbf{x}^{\prime}_i = s \cdot
+            \frac{\mathbf{x}_i^c}{{\| \mathbf{x}_i^c \|}_2}`.
+            (default: :obj:`False`)
+        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 PairNorm
+
+        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 = PairNorm()
+        out = layer(x, batch)  # normalizes each graph separately
+        print(tuple(out.shape))  # (10, 8)
+        ```
+    """
+    def __init__(self, scale: float = 1.0, scale_individually: bool = False,
+                 eps: float = 1e-5, **kwargs):
+        super().__init__(**kwargs)
+        self.scale = scale
+        self.scale_individually = scale_individually
+        self.eps = eps
+
+    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
+
+        scale = self.scale
+
+        if batch is None:
+            x = x - ops.mean(x, axis=0, keepdims=True)
+
+            if not self.scale_individually:
+                mean_sq = ops.mean(ops.sum(ops.power(x, 2), axis=-1))
+                return scale * x / ops.sqrt(self.eps + mean_sq)
+            else:
+                norm = ops.sqrt(ops.sum(ops.power(x, 2), axis=-1, keepdims=True))
+                return scale * x / (self.eps + norm)
+
+        if batch_size is not None and not isinstance(batch_size, int):
+            try:
+                batch_size = int(batch_size)
+            except Exception:
+                pass
+        elif batch_size is None:
+            batch_size = ops.cast(ops.max(batch), "int32") + 1
+
+        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
+        x = x - ops.take(mean, batch, axis=0)
+
+        if not self.scale_individually:
+            sq_sum = ops.sum(ops.power(x, 2), axis=-1, keepdims=True)
+            mean_sq = segment_sum(sq_sum, batch, num_segments=batch_size) / counts
+            denom = ops.sqrt(self.eps + ops.take(mean_sq, batch, axis=0))
+            return scale * x / denom
+        else:
+            norm = ops.sqrt(ops.sum(ops.power(x, 2), axis=-1, keepdims=True))
+            return scale * x / (self.eps + norm)
+
+    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
+