mindspore 2.1.0__cp39-cp39-win_amd64.whl → 2.2.10__cp39-cp39-win_amd64.whl

mindspore/_extends/graph_kernel/expander.py

@@ -1,80 +0,0 @@
-# Copyright 2020-2022 Huawei Technologies Co., Ltd
-#
-# Licensed 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.
-# ============================================================================
-"""generate json desc for graph kernel ops"""
-import json
-import json.decoder as jd
-import traceback
-from mindspore import log as logger
-import mindspore._extends.graph_kernel.expanders as expanders
-from mindspore._extends.graph_kernel.model.model import GraphKernelUnsupportedException
-
-
-def create_expander(expand_info):
-    """Create an expander according to op name"""
-    def call_func(func, arg):
-        return func(arg)
-    op_name = str(expand_info['name'])
-    if not hasattr(expanders, op_name):
-        raise GraphKernelUnsupportedException("Expander does not support op: {}".format(op_name))
-    expander = getattr(expanders, op_name)
-    return call_func(expander, expand_info)
-
-
-def extract_expand_info(kernel_info):
-    """Convert the json into a more friendly format"""
-    input_desc = []
-    if 'input_desc' in kernel_info and kernel_info['input_desc']:
-        for desc in kernel_info['input_desc']:
-            input_desc += desc
-    attrs = {}
-    if 'attr' in kernel_info and kernel_info['attr']:
-        for attr in kernel_info["attr"]:
-            attrs[attr["name"]] = attr["value"]
-    expand_info = {
-        "name": kernel_info["name"],
-        "input_desc": input_desc,
-        "output_desc": kernel_info["output_desc"],
-        "attr": attrs,
-        "process": kernel_info["process"],
-    }
-    return expand_info
-
-
-def get_op_expander(json_str: str):
-    """get op expander by json info"""
-    try:
-        kernel_info = json.loads(json_str)
-        expand_info = extract_expand_info(kernel_info)
-
-        expander = create_expander(expand_info)
-        graph = expander.run()
-
-        # dump graph to json desc.
-        desc = graph.dump()
-        return json.dumps(desc)
-
-    except jd.JSONDecodeError:
-        logger.error("Decode input json str failed in expander, json is: {}".format(json_str))
-        logger.error(traceback.format_exc())
-        return ""
-    except GraphKernelUnsupportedException as e:
-        logger.info(e.message)
-        return ""
-
-
-def get_expander_op_list():
-    """get supported expander op list"""
-    op_list = [name for name in dir(expanders) if name[0].isupper()]
-    return ' '.join(op_list)

mindspore/_extends/graph_kernel/expanders/__init__.py

@@ -1,54 +0,0 @@
-# Copyright 2020-2022 Huawei Technologies Co., Ltd
-#
-# Licensed 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.
-# ============================================================================
-"""expanders init. Deprecated, please add the new operators in the c++ file"""
-
-
-from .addn import AddN
-from .batchnorm import BatchNorm
-from .batchnorm_grad import BatchNormGrad
-from .clip_by_norm_no_div_sum import ClipByNormNoDivSum
-from .conv2d import Conv2D
-from .complex import CAbs, CAdd, CDiv, CMul, CSub, CRealDiv
-from .dropout_grad import DropoutGrad
-from .equal_count import EqualCount
-from .erfc import Erfc
-from .fused_adam import FusedAdam
-from .fused_adam_weight_decay import FusedAdamWeightDecay
-from .fused_mul_add import FusedMulAdd
-from .gelu_grad import GeLUGrad
-from .gkdropout import GkDropout
-from .identity import Identity
-from .layernorm import LayerNorm
-from .layernorm_grad import LayerNormGrad
-from .logsoftmax import LogSoftmax
-from .logsoftmax_grad import LogSoftmaxGrad
-from .matmul import BatchMatMul, MatMul
-from .maximum_grad import MaximumGrad
-from .minimum_grad import MinimumGrad
-from .oneslike import OnesLike
-from .reduce_mean import ReduceMean
-from .relu_grad import ReluGrad
-from .sigmoid_cross_entropy_with_logits import SigmoidCrossEntropyWithLogits
-from .sigmoid_cross_entropy_with_logits_grad import SigmoidCrossEntropyWithLogitsGrad
-from .sigmoid_grad import SigmoidGrad
-from .slice import Slice
-from .softmax_cross_entropy_with_logits import SoftmaxCrossEntropyWithLogits
-from .softmax_grad_ext import SoftmaxGradExt
-from .sqrt_grad import SqrtGrad
-from .squared_difference import SquaredDifference
-from .square_sum_v1 import SquareSumV1
-from .square_sum_all import SquareSumAll
-from .tanh_grad import TanhGrad
-from .softsign import Softsign

mindspore/_extends/graph_kernel/expanders/_utils.py

@@ -1,269 +0,0 @@
-# Copyright 2021-2022 Huawei Technologies Co., Ltd
-#
-# Licensed 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.
-# ===========================================================================
-"""GraphKernel expander utils"""
-from abc import ABCMeta, abstractmethod
-from mindspore._extends.graph_kernel.model import model_builder as builder
-from mindspore._extends.graph_kernel.model.model import GraphKernelUnsupportedException as GKException
-
-
-class Expander(metaclass=ABCMeta):
-    """
-    Expander is the base class of expanders.
-
-    The method `_expand` should be overridden to implement the operator detail.
-    """
-    def __init__(self, expand_info):
-        self.name = expand_info["name"]
-        self.inputs = expand_info["input_desc"]
-        self.outputs = expand_info["output_desc"]
-        self.attrs = expand_info["attr"]
-        self.processor = expand_info["process"]
-
-    def run(self):
-        """
-        Expand the operator to a graph.
-
-        `GraphKernelUnsupportedException` would be raised if check failed.
-        """
-        self._check()
-        graph_builder = builder.GraphBuilder()
-        with graph_builder.graph_scope(self.name) as graph_scope:
-            # transform input_desc to Tensor
-            self.inputs = [graph_builder.tensor(inp['shape'], inp['data_type'], inp['format']) for inp in self.inputs]
-            graph_scope.set_input(*self.inputs)
-            outputs = self._expand(graph_builder)
-            if isinstance(outputs, (list, tuple)):
-                self._check_output_same(outputs)
-                graph_scope.set_output(*outputs)
-            else:
-                self._check_output_same([outputs])
-                graph_scope.set_output(outputs)
-
-        graph = graph_builder.get()[0]
-        graph.set_processor(self.processor)
-        return graph
-
-    def _check(self):
-        """Check inputs"""
-
-    def _check_output_same(self, outputs):
-        for index, value in enumerate(self.outputs):
-            if list(outputs[index].shape) != list(value['shape']):
-                raise GKException("{} 's output shape {} is wrong. Expected:{}".format(
-                    self.__class__.__name__, list(outputs[index].shape), list(value['shape'])))
-            if outputs[index].dtype != value['data_type']:
-                raise GKException("{} 's output data_type {} is wrong. Expected: {}".format(
-                    self.__class__.__name__, outputs[index].dtype, value['data_type']))
-            if outputs[index].data_format != value['format']:
-                raise GKException("{} 's output format {} is wrong. Expected: {}".format(
-                    self.__class__.__name__, outputs[index].data_format, value['format']))
-
-    @abstractmethod
-    def _expand(self, graph_builder):
-        """Expand operator, this function should be overridden in subclass"""
-        raise Exception("_expand() is not implemented in {}".format(self.__class__.__name__))
-
-
-class ExpanderInfoValidator:
-    """ExpanderInfoValidator is the utility class which defines the validator decorator for expanders"""
-
-    def __init__(self):
-        """Init"""
-
-    @staticmethod
-    def _add_check_function(kls, func):
-        """
-        Rewrite the function `_check` in class Expander
-        to append the new `func` after the original checks.
-        """
-        old_check = getattr(kls, "_check")
-
-        def new_check(obj):
-            old_check(obj)
-            func(obj)
-
-        setattr(kls, "_check", new_check)
-
-    @staticmethod
-    def add_format(*input_format):
-        """
-        Add new supported format for the operator
-
-        this function will add a list `__supported_formats` into the expander,
-        saving the whitelist of formats that this op supports.
-        it also rewrites the `_check` function to check the formats.
-        """
-        format_list_name = "__supported_formats"
-
-        def _check_format(obj):
-            inp_formats = [inp['format'] for inp in obj.inputs]
-            for formats in getattr(obj, format_list_name):
-                if len(formats) != len(inp_formats):
-                    raise GKException("For '{}', length of registered format is different from the length of inputs "
-                                      "format: {} vs {}".format(obj.name, len(formats), len(inp_formats)))
-                if all((fmt == inp for fmt, inp in zip(formats, inp_formats))):
-                    return
-            raise GKException("Unregistered format ({}) for op {}".format(','.join(inp_formats), obj.name))
-
-        def wrapper(cls):
-            if not issubclass(cls, Expander):
-                raise Exception("{} should be subclass of Expander.".format(cls.__name__))
-            if not hasattr(cls, format_list_name):
-                setattr(cls, format_list_name, list())
-                ExpanderInfoValidator._add_check_function(cls, _check_format)
-            getattr(cls, format_list_name).append(input_format)
-            return cls
-
-        return wrapper
-
-    @staticmethod
-    def check_all_formats_same(kls):
-        """Check that all formats are the same"""
-
-        # Ensure no args case can return a class
-        def _check(*args):
-            def _check_format(obj):
-                inp_formats = [inp['format'] for inp in obj.inputs]
-                if all((fmt == inp_formats[0] for fmt in inp_formats[1:])):
-                    return
-                raise GKException("[check_all_formats_same] unmatched formats ({}) for op {}".format(
-                    ','.join(inp_formats), obj.name))
-
-            def wrapper(cls):
-                if not issubclass(cls, Expander):
-                    raise Exception("{} should be subclass of Expander.".format(cls.__name__))
-                ExpanderInfoValidator._add_check_function(cls, _check_format)
-                return cls
-
-            return wrapper
-
-        return _check()(kls)
-
-    @staticmethod
-    def check_attrs(*args):
-        """Check the attrs exist"""
-
-        def _check_attr(obj):
-            for a in args:
-                if a not in obj.attrs:
-                    raise GKException("attr '{}' does not exist. {}".format(a, obj.name))
-
-        def wrapper(cls):
-            if not issubclass(cls, Expander):
-                raise Exception("{} should be subclass of Expander.".format(cls.__name__))
-            ExpanderInfoValidator._add_check_function(cls, _check_attr)
-            return cls
-
-        return wrapper
-
-
-def to_frac_z_axis(ori_shape, ori_axis):
-    """
-    judge the format is fractal NZ
-    Parameters
-    ----------
-    ori_shape: list or tuple
-        original shape of input
-    ori_axis: list or tuple
-        original axis of original shape to operate
-    Returns
-    -------
-    output: list
-        axis of the fractal Nz shape
-    """
-    frac_z_axis = list(ori_axis)
-    shape_len = len(ori_shape)
-    axis_count = len(frac_z_axis)
-    axis_negative_1 = shape_len - 1
-    axis_negative_2 = shape_len - 2
-    for i in range(axis_count):
-        axis_index = (frac_z_axis[i] + shape_len) % shape_len
-        if axis_index == axis_negative_1:
-            if frac_z_axis[i] > shape_len - 2:  # akg:[2,3] [1,4] tbe:[2,4] [1,3]
-                frac_z_axis[i] = axis_index - 1
-                frac_z_axis.append(axis_index + 2)
-            else:  # no case cover this branch now
-                frac_z_axis[i] = axis_index - 1
-                frac_z_axis.append(axis_index + 2)
-        elif axis_index == axis_negative_2:
-            frac_z_axis[i] = axis_index + 1
-            frac_z_axis.append(axis_index + 2)
-        else:
-            frac_z_axis[i] = axis_index
-    return frac_z_axis
-
-
-def infer_shape_from_fractalnz(fractal):
-    "get original shape from fractalnz shape"
-    shape = []
-    dims = len(fractal)
-    batch = dims - 4
-    for i in range(batch):
-        shape.append(fractal[i])
-    m = fractal[dims - 3] * fractal[dims - 2]
-    n = fractal[dims - 4] * fractal[dims - 1]
-    shape.append(m)
-    shape.append(n)
-    return shape
-
-
-def get_reduced_ori_shape(shape, axis):
-    "get shape after reduced which is based on original shape"
-    reduced_ori_shape = []
-    for i, value in enumerate(shape):
-        if i in axis:
-            reduced_ori_shape.append(1)
-        else:
-            reduced_ori_shape.append(value)
-    return reduced_ori_shape
-
-
-def get_reduce_axis_shape(shape, data_format, axis):
-    """
-    Get the reduce axis under format `data_format` and original reduced shape.
-    Parameters
-    ----------
-    shape: list or tuple
-        shape of input
-    data_format: str
-        data format of input
-    axis: None, int, list or tuple
-        reduce axis of the original shape
-    Returns
-    -------
-    reduce_axis: list
-        reduce axis of the `data_format` shape
-    ori_reduced_shape: list
-        original reduced shape
-    """
-    ori_shape = shape
-    if data_format == "FRACTAL_NZ":
-        ori_shape = infer_shape_from_fractalnz(shape)
-    if not axis:
-        axis = []
-        for i, _ in enumerate(ori_shape):
-            axis.append(i)
-    else:
-        if isinstance(axis, int):
-            axis = [axis]
-        for i, _ in enumerate(list(axis)):
-            if axis[i] < 0:
-                axis[i] += len(ori_shape)
-
-    ori_reduced_shape = get_reduced_ori_shape(ori_shape, axis)
-    reduce_axis = axis
-    if data_format == "FRACTAL_NZ":
-        reduce_axis = to_frac_z_axis(ori_shape, axis)
-    return reduce_axis, ori_reduced_shape

mindspore/_extends/graph_kernel/expanders/addn.py

@@ -1,33 +0,0 @@
-# Copyright 2021-2022 Huawei Technologies Co., Ltd
-#
-# Licensed 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.
-# ===========================================================================
-"""generate json desc for addn"""
-from mindspore._extends.graph_kernel.model.model import GraphKernelUnsupportedException as GKException
-from ._utils import Expander, ExpanderInfoValidator as VLD
-
-
-@VLD.check_all_formats_same
-class AddN(Expander):
-    """Expand AddN to multiple Adds"""
-
-    def _check(self):
-        if len(self.inputs) < 2:
-            raise GKException("For 'AddN', the inputs num should be greater than 1, but got {}"
-                              .format(len(self.inputs)))
-
-    def _expand(self, graph_builder):
-        result = self.inputs[0]
-        for inp in self.inputs[1:]:
-            result = graph_builder.emit('Add', [result, inp])
-        return result

mindspore/_extends/graph_kernel/expanders/batchnorm.py

@@ -1,152 +0,0 @@
-# Copyright 2021-2022 Huawei Technologies Co., Ltd
-#
-# Licensed 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.
-# ===========================================================================
-"""generate json desc for BatchNorm"""
-from mindspore._extends.graph_kernel.model.model import DataFormat as DF
-from ._utils import Expander, ExpanderInfoValidator as VLD
-from .expand_dims import ExpandDims
-
-
-@VLD.add_format(DF.NHWC, DF.DEFAULT, DF.DEFAULT, DF.DEFAULT, DF.DEFAULT)
-@VLD.add_format(DF.NCHW, DF.DEFAULT, DF.DEFAULT, DF.DEFAULT, DF.DEFAULT)
-@VLD.add_format(DF.DEFAULT, DF.DEFAULT, DF.DEFAULT, DF.DEFAULT, DF.DEFAULT)
-@VLD.check_attrs('is_training', 'momentum', 'epsilon')
-class BatchNorm(Expander):
-    """BatchNorm expander"""
-
-    def _expand(self, graph_builder):
-        # get op info
-        input_x = self.inputs[0]
-        input_scale = self.inputs[1]
-        input_offset = self.inputs[2]
-        input_mean = self.inputs[3]
-        input_variance = self.inputs[4]
-        epsilon_v = graph_builder.value(input_scale.dtype, self.attrs['epsilon'])
-
-        input_x_ori_type = input_x.dtype
-        input_x_new_type = input_x.dtype
-        if input_x.dtype == "float16" and input_scale.dtype == "float32" and input_offset.dtype == "float32" and \
-                input_mean.dtype == "float32" and input_variance.dtype == "float32":
-            input_x_new_type = "float32"
-        if input_x_new_type != input_x_ori_type:
-            input_x = graph_builder.emit('Cast', [input_x], attrs={'dst_type': input_x_new_type})
-
-        if self.attrs['is_training']:
-            self.inputs[0] = input_x
-            res_y, mean_res, variance_res, mean_muls, y_sqrt_rec = self._bn_train(graph_builder)
-            if input_x_new_type != input_x_ori_type:
-                res_y = graph_builder.emit('Cast', [res_y], attrs={'dst_type': input_x_ori_type})
-            return res_y, mean_res, variance_res, mean_muls, y_sqrt_rec
-        # infer mode
-        if input_x.data_format in (DF.DEFAULT, DF.NCHW):
-            input_mean = graph_builder.emit(
-                'Reshape', [input_mean], attrs={'shape': ExpandDims.infer_shape(input_mean.shape, [-1, -1])})
-            input_scale = graph_builder.emit(
-                'Reshape', [input_scale], attrs={'shape': ExpandDims.infer_shape(input_scale.shape, [-1, -1])})
-            input_offset = graph_builder.emit(
-                'Reshape', [input_offset], attrs={'shape': ExpandDims.infer_shape(input_offset.shape, [-1, -1])})
-        x_sub = graph_builder.emit('Sub', [input_x, input_mean])
-        x_sub_mul = graph_builder.emit('Mul', [input_scale, x_sub])
-        var_add = graph_builder.emit('Add', [epsilon_v, input_variance])
-        var_add_sqrt = graph_builder.emit('Sqrt', [var_add])
-        if input_x.data_format in (DF.DEFAULT, DF.NCHW):
-            var_add_sqrt = graph_builder.emit(
-                'Reshape', [var_add_sqrt], attrs={'shape': ExpandDims.infer_shape(var_add_sqrt.shape, [-1, -1])})
-        x_div = graph_builder.emit('RealDiv', [x_sub_mul, var_add_sqrt])
-        res_y = graph_builder.emit('Add', [input_offset, x_div])
-        if input_x_new_type != input_x_ori_type:
-            res_y = graph_builder.emit('Cast', [res_y], attrs={'dst_type': input_x_ori_type})
-        return res_y, var_add, var_add, var_add, var_add
-
-    def _bn_train(self, graph_builder):
-        """expand BatchNorm for training mode"""
-        input_x = self.inputs[0]
-        input_scale = self.inputs[1]
-        input_offset = self.inputs[2]
-        input_mean = self.inputs[3]
-        input_variance = self.inputs[4]
-        epsilon_v = graph_builder.value(input_scale.dtype, self.attrs['epsilon'])
-        reduce_axis = ()
-        shape_x = input_x.shape
-        if input_x.data_format == DF.NHWC:
-            reduce_axis = (0, 1, 2)
-            num = shape_x[0] * shape_x[1] * shape_x[2]
-        else:
-            reduce_axis = (0, 2, 3)
-            num = shape_x[0] * shape_x[2] * shape_x[3]
-        num_rec = 1.0 / num
-        num_rec_v = graph_builder.value(input_scale.dtype, num_rec)
-
-        # compute mean value of input_x
-        mean_sum = graph_builder.emit(
-            'ReduceSum', [input_x], attrs={'reduce_axis': reduce_axis, 'keep_dims': False})
-        mean_muls = graph_builder.emit('Mul', [mean_sum, num_rec_v])
-
-        # compute variance of input_x
-        if input_x.data_format in (DF.DEFAULT, DF.NCHW):
-            mean_muls_expand = graph_builder.emit(
-                'Reshape', [mean_muls], attrs={'shape': ExpandDims.infer_shape(mean_muls.shape, [-1, -1])})
-        else:
-            mean_muls_expand = mean_muls
-        var_sub = graph_builder.emit('Sub', [input_x, mean_muls_expand])
-        var_mul = graph_builder.emit('Mul', [var_sub, var_sub])
-        var_sum = graph_builder.emit('ReduceSum', [var_mul], attrs={'reduce_axis': reduce_axis, 'keep_dims': False})
-        var_mul = graph_builder.emit('Mul', [var_sum, num_rec_v])
-
-        # y_sqrt_rec means 1 / sqrt(variance + epsilon), which is calculated in backward pass
-        scalar_one = 1.0
-        scalar_one_v = graph_builder.value(input_scale.dtype, scalar_one)
-        y_add = graph_builder.emit('Add', [var_mul, epsilon_v])
-        y_sqrt = graph_builder.emit('Sqrt', [y_add])
-        y_sqrt_rec = graph_builder.emit('RealDiv', [scalar_one_v, y_sqrt])
-
-        # compute res_y
-        tmp_sub = graph_builder.emit('Sub', [input_x, mean_muls_expand])
-        if input_x.data_format in (DF.DEFAULT, DF.NCHW):
-            y_sqrt_rec_expand = graph_builder.emit(
-                'Reshape', [y_sqrt_rec], attrs={'shape': ExpandDims.infer_shape(y_sqrt_rec.shape, [-1, -1])})
-        else:
-            y_sqrt_rec_expand = y_sqrt_rec
-        y_norm = graph_builder.emit('Mul', [tmp_sub, y_sqrt_rec_expand])
-        if input_x.data_format in (DF.DEFAULT, DF.NCHW):
-            input_scale_expand = graph_builder.emit(
-                'Reshape', [input_scale], attrs={'shape': ExpandDims.infer_shape(input_scale.shape, [-1, -1])})
-        else:
-            input_scale_expand = input_scale
-        res_y_mul = graph_builder.emit('Mul', [input_scale_expand, y_norm])
-        if input_x.data_format in (DF.DEFAULT, DF.NCHW):
-            input_offset_expand = graph_builder.emit(
-                'Reshape', [input_offset], attrs={'shape': ExpandDims.infer_shape(input_offset.shape, [-1, -1])})
-        else:
-            input_offset_expand = input_offset
-        res_y = graph_builder.emit('Add', [res_y_mul, input_offset_expand])
-
-        # compute mean_res
-        momentum_sub = scalar_one - self.attrs['momentum']
-        momentum_v_sub = graph_builder.value(input_scale.dtype, momentum_sub)
-        new_running_mean_tmp = graph_builder.emit('Mul', [momentum_v_sub, input_mean])
-        momentum_v = graph_builder.value(input_scale.dtype, self.attrs['momentum'])
-        current_mean_tmp = graph_builder.emit('Mul', [momentum_v, mean_muls])
-        updated_moving_mean = graph_builder.emit('Add', [new_running_mean_tmp, current_mean_tmp])
-        mean_res = graph_builder.emit('Assign', [input_mean, updated_moving_mean])
-
-        # variance_res is calculated by sample variance, and need to multiply by num / (num - 1)
-        var_num = float(num) / (num - 1)
-        var_num_v = graph_builder.value(input_scale.dtype, var_num)
-        var_mul_update = graph_builder.emit('Mul', [var_num_v, var_mul])
-        new_running_var_tmp = graph_builder.emit('Mul', [momentum_v_sub, input_variance])
-        current_var_tmp = graph_builder.emit('Mul', [momentum_v, var_mul_update])
-        updated_moving_variance = graph_builder.emit('Add', [new_running_var_tmp, current_var_tmp])
-        variance_res = graph_builder.emit('Assign', [input_variance, updated_moving_variance])
-        return res_y, mean_res, variance_res, mean_muls, y_sqrt_rec