chore: import upstream snapshot with attribution

2026-07-13 13:36:25 +08:00
commit 26446540fa
3151 changed files with 974126 additions and 0 deletions
@@ -0,0 +1,146 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Batch normalization."""
+
+from functools import reduce
+
+from tvm import te, topi
+
+
+def batch_norm(
+    data: te.Tensor,
+    gamma: te.Tensor,
+    beta: te.Tensor,
+    moving_mean: te.Tensor,
+    moving_var: te.Tensor,
+    axis: int | None = None,
+    epsilon: float | None = None,
+    center: bool | None = None,
+    scale: bool | None = None,
+    training: bool | None = None,
+    momentum: float | None = None,
+) -> list[te.Tensor]:
+    """Batch normalization layer (Ioffe and Szegedy, 2014).
+
+    Normalizes the input at each batch, i.e. applies a transformation
+    that maintains the mean activation close to 0 and the activation
+    standard deviation close to 1.
+
+    Parameters
+    ----------
+    data : tvm.te.Tensor
+        Input to be batch-normalized.
+
+    gamma : tvm.te.Tensor
+        Scale factor to be applied to the normalized tensor.
+
+    beta : tvm.te.Tensor
+        Offset to be applied to the normalized tensor.
+
+    moving_mean : tvm.te.Tensor
+        Running mean of input.
+
+    moving_var : tvm.te.Tensor
+        Running variance of input.
+
+    axis : int, optional, default=1
+        Specify along which shape axis the normalization should occur.
+
+    epsilon : float, optional, default=1e-5
+        Small float added to variance to avoid dividing by zero.
+
+    center : bool, optional, default=True
+        If True, add offset of beta to normalized tensor, If False,
+        beta is ignored.
+
+    scale : bool, optional, defualt=True
+        If True, scale normalized tensor by gamma. If False, gamma
+        is ignored.
+
+    training : bool, optional, defualt=False
+        Indicating whether it is in training mode. If True, update
+        moving_mean and moving_var.
+
+    momentum : float, optional, default=0.1
+        The value used for the moving_mean and moving_var update.
+
+    Returns
+    -------
+    output : list of tvm.te.Tensor
+        Normalized data with same shape as input
+
+    moving_mean : tvm.te.Tensor
+        Running mean of input.
+
+    moving_var : tvm.te.Tensor
+        Running variance of input.
+    """
+    if axis is None:
+        axis = 1
+
+    if epsilon is None:
+        epsilon = 1e-5
+
+    if center is None:
+        center = True
+
+    if scale is None:
+        scale = True
+
+    if training is None:
+        training = False
+
+    if momentum is None:
+        momentum = 0.1
+
+    shape = [1] * len(data.shape)
+    shape[axis] = data.shape[axis]
+    data_mean = None
+    data_var = None
+
+    if training:
+        reduce_axes = list(range(len(data.shape)))
+        reduce_axes.remove(axis)
+        shape_prod = reduce(lambda x, y: x * y, [data.shape[ax] for ax in reduce_axes], 1)
+        data_mean = topi.sum(data, axis=reduce_axes) / shape_prod
+        data_mean_rs = topi.reshape(data_mean, shape)
+        data_var = (
+            topi.sum((data - data_mean_rs) * (data - data_mean_rs), axis=reduce_axes) / shape_prod
+        )
+        data_var_rs = topi.reshape(data_var, shape)
+        out = (data - data_mean_rs) / topi.math.sqrt(data_var_rs + epsilon)
+    else:
+        moving_mean_rs = topi.reshape(moving_mean, shape)
+        moving_var_rs = topi.reshape(moving_var, shape)
+        out = (data - moving_mean_rs) / topi.math.sqrt(moving_var_rs + epsilon)
+
+    if scale:
+        out = out * topi.reshape(gamma, shape)
+    if center:
+        out = out + topi.reshape(beta, shape)
+
+    if training:
+        assert 0 <= momentum <= 1, "the valid momentum range is [0, 1]."
+        return [
+            out,
+            (1 - momentum) * moving_mean + momentum * data_mean,
+            (1 - momentum) * moving_var + momentum * data_var,
+        ]
+
+    # Moving mean and var aren't updated during test. To avoid
+    # placeholder reuse, we multiply by 1 and return them.
+    return [out, moving_mean * 1, moving_var * 1]