mindspore 1.10.0__cp37-cp37m-win_amd64.whl → 2.0.0rc1__cp37-cp37m-win_amd64.whl

mindspore/ops/operations/sponge_ops.py

@@ -1,3531 +0,0 @@
-# Copyright 2021 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.
-# ============================================================================
-
-"""
-Note:
-  SPONGE operators. This is an experimental interface that is subject to change and/or deletion.
-"""
-
-import math
-
-from ..primitive import PrimitiveWithInfer, prim_attr_register
-from ..._checkparam import Rel
-from ..._checkparam import Validator as validator
-from ...common import dtype as mstype
-
-
-class BondForce(PrimitiveWithInfer):
-    """
-    Calculate the force exerted by the simple harmonic bond on the corresponding atoms.
-    Assume the number of harmonic bonds is m and the number of atoms is n.
-
-    Because there is a large amount of inputs and each of them are related,
-    there is no way to construct `Examples` using random methods. For details, refer the webpage `SPONGE in MindSpore
-    <https://gitee.com/mindspore/mindscience/blob/master/MindSPONGE/docs/simple_formula.md>`_.
-
-    .. math::
-
-        dr = (x_1-x_2, y_1-y_2, z_1-z_2)
-
-    .. math::
-
-        F = (F_x, F_y, F_z) = 2*k*(1 - r_0/|dr|)*dr
-
-    Args:
-        atom_numbers(int32): the number of atoms n.
-        bond_numbers(int32): the number of harmonic bonds m.
-
-    Inputs:
-        - **uint_crd_f** (Tensor) - The unsigned int coordinate value of each atom.
-          The data type is uint32 and the shape is :math:`(n, 3)`.
-        - **scaler_f** (Tensor) - The 3-D scale factor (x, y, z),
-          between the real space float coordinates and the unsigned int coordinates.
-          The data type is float32  and the shape is :math:`(3,)`.
-        - **atom_a** (Tensor) - The first atom index of each bond.
-          The data type is int32 and the shape is :math:`(m,)`.
-        - **atom_b** (Tensor) - The second atom index of each bond.
-          The data type is int32 and the shape is :math:`(m,)`.
-        - **bond_k** (Tensor) - The force constant of each bond.
-          The data type is float32 and the shape is :math:`(m,)`.
-        - **bond_r0** (Tensor) - The equlibrium length of each bond.
-          The data type is float32 and the shape is :math:`(m,)`.
-
-    Outputs:
-        - **frc_f** (Tensor) - The force felt by each atom.
-          The data type is float32 and the shape is :math:`(n, 3)`.
-
-    Supported Platforms:
-        ``GPU``
-    """
-
-    @prim_attr_register
-    def __init__(self, bond_numbers, atom_numbers):
-        """Initialize BondForce."""
-        validator.check_value_type('bond_numbers', bond_numbers, int, self.name)
-        validator.check_value_type('atom_numbers', atom_numbers, int, self.name)
-        self.bond_numbers = bond_numbers
-        self.atom_numbers = atom_numbers
-        self.add_prim_attr('bond_numbers', self.bond_numbers)
-        self.add_prim_attr('atom_numbers', self.atom_numbers)
-        self.init_prim_io_names(inputs=['uint_crd_f', 'scaler_f', 'atom_a', 'atom_b', 'bond_k', 'bond_r0'],
-                                outputs=['frc_f'])
-
-    def infer_shape(self, uint_crd_f_shape, scaler_f_shape, atom_a_shape, atom_b_shape, bond_k_shape, bond_r0_shape):
-        cls_name = self.name
-        n = self.atom_numbers
-        m = self.bond_numbers
-        validator.check_int(len(uint_crd_f_shape), 2, Rel.EQ, "uint_crd_f_dim", cls_name)
-        validator.check_int(len(scaler_f_shape), 1, Rel.EQ, "scaler_f_dim", cls_name)
-        validator.check_int(len(atom_a_shape), 1, Rel.EQ, "atom_a_dim", cls_name)
-        validator.check_int(len(atom_b_shape), 1, Rel.EQ, "atom_b_dim", cls_name)
-        validator.check_int(len(bond_k_shape), 1, Rel.EQ, "bond_k_dim", cls_name)
-        validator.check_int(len(bond_r0_shape), 1, Rel.EQ, "bond_r0_dim", cls_name)
-        validator.check_int(uint_crd_f_shape[0], n, Rel.EQ, "uint_crd_f_shape[0]", cls_name)
-        validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape[1]", cls_name)
-        validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f_shape", cls_name)
-        validator.check_int(atom_a_shape[0], m, Rel.EQ, "uint_crd_f_shape", cls_name)
-        validator.check_int(atom_b_shape[0], m, Rel.EQ, "atom_b_shape", cls_name)
-        validator.check_int(bond_k_shape[0], m, Rel.EQ, "bond_k_shape", cls_name)
-        validator.check_int(bond_r0_shape[0], m, Rel.EQ, "bond_r0_shape", cls_name)
-        return uint_crd_f_shape
-
-    def infer_dtype(self, uint_crd_f_dtype, scaler_f_type, atom_a_type, atom_b_type, bond_k_type, bond_r0_type):
-        validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name)
-        validator.check_tensor_dtype_valid('scaler_f', scaler_f_type, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('atom_a', atom_a_type, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('atom_b', atom_b_type, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('bond_k', bond_k_type, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('bond_r0', bond_r0_type, [mstype.float32], self.name)
-        return bond_r0_type
-
-
-class BondEnergy(PrimitiveWithInfer):
-    """
-    Calculate the harmonic potential energy between each bonded atom pair.
-    Assume our system has n atoms and m harmonic bonds.
-
-    Because there is a large amount of inputs and each of them are related,
-    there is no way to construct `Examples` using random methods. For details, refer the webpage `SPONGE in MindSpore
-    <https://gitee.com/mindspore/mindscience/blob/master/MindSPONGE/docs/simple_formula.md>`_.
-
-    .. math::
-
-        dr = (x_1-x_2, y_1-y_2, z_1-z_2)
-
-    .. math::
-
-        E = k*(|dr| - r_0)^2
-
-    Args:
-        atom_numbers(int32): the number of atoms n.
-        bond_numbers(int32): the number of harmonic bonds m.
-
-    Inputs:
-        - **uint_crd_f** (Tensor) - The unsigned int coordinate value of each atom.
-          The data type is uint32 and the shape is :math:`(n, 3)`.
-        - **scaler_f** (Tensor) - The 3-D scale factor (x, y, z),
-          between the real space float coordinates and the unsigned int coordinates.
-          The data type is float32 and the shape is :math:`(3,)`.
-        - **atom_a** (Tensor) - The first atom index of each bond.
-          The data type is int32 and the shape is :math:`(m,)`.
-        - **atom_b** (Tensor) - The second atom index of each bond.
-          The data type is int32 and the shape is :math:`(m,)`.
-        - **bond_k** (Tensor) - The force constant of each bond.
-          The data type is float32 and the shape is :math:`(m,)`.
-        - **bond_r0** (Tensor) - The equlibrium length of each bond.
-          The data type is float32 and the shape is :math:`(m,)`.
-
-    Outputs:
-        - **bond_ene** (Tensor) - The harmonic potential energy for each bond.
-          The data type is float32 and the shape is :math:`(m,)`.
-
-    Supported Platforms:
-        ``GPU``
-    """
-
-    @prim_attr_register
-    def __init__(self, bond_numbers, atom_numbers):
-        """Initialize BondEnergy."""
-        validator.check_value_type('bond_numbers', bond_numbers, int, self.name)
-        validator.check_value_type('atom_numbers', atom_numbers, int, self.name)
-        self.bond_numbers = bond_numbers
-        self.atom_numbers = atom_numbers
-        self.add_prim_attr('bond_numbers', self.bond_numbers)
-        self.add_prim_attr('atom_numbers', self.atom_numbers)
-        self.init_prim_io_names(inputs=['uint_crd_f', 'scaler_f', 'atom_a', 'atom_b', 'bond_k', 'bond_r0'],
-                                outputs=['bond_ene'])
-
-    def infer_shape(self, uint_crd_f_shape, scaler_f_shape, atom_a_shape, atom_b_shape, bond_k_shape, bond_r0_shape):
-        cls_name = self.name
-        n = self.atom_numbers
-        m = self.bond_numbers
-        validator.check_int(len(uint_crd_f_shape), 2, Rel.EQ, "uint_crd_f_dim", cls_name)
-        validator.check_int(len(scaler_f_shape), 1, Rel.EQ, "scaler_f_dim", cls_name)
-        validator.check_int(len(atom_a_shape), 1, Rel.EQ, "atom_a_dim", cls_name)
-        validator.check_int(len(atom_b_shape), 1, Rel.EQ, "atom_b_dim", cls_name)
-        validator.check_int(len(bond_k_shape), 1, Rel.EQ, "bond_k_dim", cls_name)
-        validator.check_int(len(bond_r0_shape), 1, Rel.EQ, "bond_r0_dim", cls_name)
-        validator.check_int(uint_crd_f_shape[0], n, Rel.EQ, "uint_crd_f_shape[0]", cls_name)
-        validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape[1]", cls_name)
-        validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f_shape", cls_name)
-        validator.check_int(atom_a_shape[0], m, Rel.EQ, "uint_crd_f_shape", cls_name)
-        validator.check_int(atom_b_shape[0], m, Rel.EQ, "atom_b_shape", cls_name)
-        validator.check_int(bond_k_shape[0], m, Rel.EQ, "bond_k_shape", cls_name)
-        validator.check_int(bond_r0_shape[0], m, Rel.EQ, "bond_r0_shape", cls_name)
-
-        return bond_k_shape
-
-    def infer_dtype(self, uint_crd_f_dtype, scaler_f_type, atom_a_type, atom_b_type, bond_k_type, bond_r0_type):
-        validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name)
-        validator.check_tensor_dtype_valid('scaler_f', scaler_f_type, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('atom_a', atom_a_type, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('atom_b', atom_b_type, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('bond_k', bond_k_type, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('bond_r0', bond_r0_type, [mstype.float32], self.name)
-        return bond_r0_type
-
-
-class BondAtomEnergy(PrimitiveWithInfer):
-    """
-    Add the potential energy caused by simple harmonic bonds to the total
-    potential energy of each atom.
-
-    The calculation formula is the same as operator BondEnergy().
-
-    Because there is a large amount of inputs and each of them are related,
-    there is no way to construct `Examples` using random methods. For details, refer the webpage `SPONGE in MindSpore
-    <https://gitee.com/mindspore/mindscience/blob/master/MindSPONGE/docs/simple_formula.md>`_.
-
-    Args:
-        atom_numbers(int32): the number of atoms n.
-        bond_numbers(int32): the number of harmonic bonds m.
-
-    Inputs:
-        - **uint_crd_f** (Tensor) - The unsigned int coordinate value of each atom.
-          The data type is uint32 and the shape is :math:`(n, 3)`.
-        - **scaler_f** (Tensor) - The 3-D scale factor (x, y, z),
-          between the real space float coordinates and the unsigned int coordinates.
-          The data type is float32 and the shape is :math:`(3,)`.
-        - **atom_a** (Tensor) - The first atom index of each bond.
-          The data type is int32 and the shape is :math:`(m,)`.
-        - **atom_b** (Tensor) - The second atom index of each bond.
-          The data type is int32 and the shape is :math:`(m,)`.
-        - **bond_k** (Tensor) - The force constant of each bond.
-          The data type is float32 and the shape is :math:`(m,)`.
-        - **bond_r0** (Tensor) - The equlibrium length of each bond.
-          The data type is float32 and the shape is :math:`(m,)`.
-
-    Outputs:
-        - **atom_ene** (Tensor) - The accumulated potential energy for each atom.
-          The data type is float32 and the shape is :math:`(n,)`.
-
-    Supported Platforms:
-        ``GPU``
-    """
-
-    @prim_attr_register
-    def __init__(self, bond_numbers, atom_numbers):
-        """Initialize BondAtomEnergy."""
-        validator.check_value_type('bond_numbers', bond_numbers, int, self.name)
-        validator.check_value_type('atom_numbers', atom_numbers, int, self.name)
-        self.bond_numbers = bond_numbers
-        self.atom_numbers = atom_numbers
-        self.add_prim_attr('bond_numbers', self.bond_numbers)
-        self.add_prim_attr('atom_numbers', self.atom_numbers)
-        self.init_prim_io_names(inputs=['uint_crd_f', 'scaler_f', 'atom_a', 'atom_b', 'bond_k', 'bond_r0'],
-                                outputs=['atom_ene'])
-
-    def infer_shape(self, uint_crd_f_shape, scaler_f_shape, atom_a_shape, atom_b_shape, bond_k_shape, bond_r0_shape):
-        cls_name = self.name
-        n = self.atom_numbers
-        m = self.bond_numbers
-        validator.check_int(len(uint_crd_f_shape), 2, Rel.EQ, "uint_crd_f_dim", cls_name)
-        validator.check_int(len(scaler_f_shape), 1, Rel.EQ, "scaler_f_dim", cls_name)
-        validator.check_int(len(atom_a_shape), 1, Rel.EQ, "atom_a_dim", cls_name)
-        validator.check_int(len(atom_b_shape), 1, Rel.EQ, "atom_b_dim", cls_name)
-        validator.check_int(len(bond_k_shape), 1, Rel.EQ, "bond_k_dim", cls_name)
-        validator.check_int(len(bond_r0_shape), 1, Rel.EQ, "bond_r0_dim", cls_name)
-        validator.check_int(uint_crd_f_shape[0], n, Rel.EQ, "uint_crd_f_shape[0]", cls_name)
-        validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape[1]", cls_name)
-        validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f_shape", cls_name)
-        validator.check_int(atom_a_shape[0], m, Rel.EQ, "uint_crd_f_shape", cls_name)
-        validator.check_int(atom_b_shape[0], m, Rel.EQ, "atom_b_shape", cls_name)
-        validator.check_int(bond_k_shape[0], m, Rel.EQ, "bond_k_shape", cls_name)
-        validator.check_int(bond_r0_shape[0], m, Rel.EQ, "bond_r0_shape", cls_name)
-        return [n,]
-
-    def infer_dtype(self, uint_crd_f_dtype, scaler_f_type, atom_a_type, atom_b_type, bond_k_type, bond_r0_type):
-        validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name)
-        validator.check_tensor_dtype_valid('scaler_f', scaler_f_type, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('atom_a', atom_a_type, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('atom_b', atom_b_type, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('bond_k', bond_k_type, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('bond_r0', bond_r0_type, [mstype.float32], self.name)
-        return bond_r0_type
-
-
-class BondForceWithAtomEnergy(PrimitiveWithInfer):
-    """
-    Calculate bond force and harmonic potential energy together.
-
-    The calculation formula is the same as operator BondForce() and BondEnergy().
-
-    Because there is a large amount of inputs and each of them are related,
-    there is no way to construct `Examples` using random methods. For details, refer the webpage `SPONGE in MindSpore
-    <https://gitee.com/mindspore/mindscience/blob/master/MindSPONGE/docs/simple_formula.md>`_.
-
-    Args:
-        atom_numbers(int32): the number of atoms n.
-        bond_numbers(int32): the number of harmonic bonds m.
-
-    Inputs:
-        - **uint_crd_f** (Tensor) - The unsigned int coordinate value of each atom.
-          The data type is uint32 and the shape is :math:`(n, 3)`.
-        - **scaler_f** (Tensor) - The 3-D scale factor (x, y, z),
-          between the real space float coordinates and the unsigned int coordinates.
-          The data type is float32 and the shape is :math:`(3,)`.
-        - **atom_a** (Tensor) - The first atom index of each bond.
-          The data type is int32 and the shape is :math:`(m,)`.
-        - **atom_b** (Tensor) - The second atom index of each bond.
-          The data type is int32 and the shape is :math:`(m,)`.
-        - **bond_k** (Tensor) - The force constant of each bond.
-          The data type is float32 and the shape is :math:`(m,)`.
-        - **bond_r0** (Tensor) - The equlibrium length of each bond.
-          The data type is float32 and the shape is :math:`(m,)`.
-
-    Outputs:
-        - **frc_f** (Tensor) - The force felt by each atom.
-          The data type is float32 and the shape is :math:`(n, 3)`.
-        - **atom_e** (Tensor) - The accumulated potential energy for each atom.
-          The data type is float32 and the shape is :math:`(n,)`.
-
-    Supported Platforms:
-        ``GPU``
-    """
-
-    @prim_attr_register
-    def __init__(self, bond_numbers, atom_numbers):
-        """Initialize BondForceWithAtomEnergy."""
-        validator.check_value_type('bond_numbers', bond_numbers, int, self.name)
-        validator.check_value_type('atom_numbers', atom_numbers, int, self.name)
-        self.bond_numbers = bond_numbers
-        self.atom_numbers = atom_numbers
-        self.add_prim_attr('bond_numbers', self.bond_numbers)
-        self.add_prim_attr('atom_numbers', self.atom_numbers)
-        self.init_prim_io_names(inputs=['uint_crd_f', 'scaler_f', 'atom_a', 'atom_b', 'bond_k', 'bond_r0'],
-                                outputs=['frc_f', 'atom_e'])
-
-    def infer_shape(self, uint_crd_f_shape, scaler_f_shape, atom_a_shape, atom_b_shape, bond_k_shape, bond_r0_shape):
-        cls_name = self.name
-        n = self.atom_numbers
-        m = self.bond_numbers
-        validator.check_int(len(uint_crd_f_shape), 2, Rel.EQ, "uint_crd_f_dim", cls_name)
-        validator.check_int(len(scaler_f_shape), 1, Rel.EQ, "scaler_f_dim", cls_name)
-        validator.check_int(len(atom_a_shape), 1, Rel.EQ, "atom_a_dim", cls_name)
-        validator.check_int(len(atom_b_shape), 1, Rel.EQ, "atom_b_dim", cls_name)
-        validator.check_int(len(bond_k_shape), 1, Rel.EQ, "bond_k_dim", cls_name)
-        validator.check_int(len(bond_r0_shape), 1, Rel.EQ, "bond_r0_dim", cls_name)
-        validator.check_int(uint_crd_f_shape[0], n, Rel.EQ, "uint_crd_f_shape[0]", cls_name)
-        validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape[1]", cls_name)
-        validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f_shape", cls_name)
-        validator.check_int(atom_a_shape[0], m, Rel.EQ, "uint_crd_f_shape", cls_name)
-        validator.check_int(atom_b_shape[0], m, Rel.EQ, "atom_b_shape", cls_name)
-        validator.check_int(bond_k_shape[0], m, Rel.EQ, "bond_k_shape", cls_name)
-        validator.check_int(bond_r0_shape[0], m, Rel.EQ, "bond_r0_shape", cls_name)
-
-        return uint_crd_f_shape, [n,]
-
-    def infer_dtype(self, uint_crd_f_dtype, scaler_f_type, atom_a_type, atom_b_type, bond_k_type, bond_r0_type):
-        validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name)
-        validator.check_tensor_dtype_valid('scaler_f', scaler_f_type, [mstype.float32], self.name)
-
-        validator.check_tensor_dtype_valid('atom_a', atom_a_type, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('atom_b', atom_b_type, [mstype.int32], self.name)
-
-        validator.check_tensor_dtype_valid('bond_k', bond_k_type, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('bond_r0', bond_r0_type, [mstype.float32], self.name)
-        return bond_r0_type, bond_r0_type
-
-
-class BondForceWithAtomVirial(PrimitiveWithInfer):
-    """
-    Calculate bond force and the virial coefficient caused by simple harmonic
-    bond for each atom together.
-
-    The calculation formula of the force part is the same as operator BondForce().
-
-    Because there is a large amount of inputs and each of them are related,
-    there is no way to construct `Examples` using random methods. For details, refer the webpage `SPONGE in MindSpore
-    <https://gitee.com/mindspore/mindscience/blob/master/MindSPONGE/docs/simple_formula.md>`_.
-
-    The Virial part is as follows:
-
-    .. math::
-
-        dr = (x_1-x_2, y_1-y_2, z_1-z_2)
-
-    .. math::
-
-        virial = |dr|*(|dr| - r_0)*k
-
-    Args:
-        atom_numbers(int32): the number of atoms n.
-        bond_numbers(int32): the number of harmonic bonds m.
-
-    Inputs:
-        - **uint_crd_f** (Tensor) - The unsigned int coordinate value of each atom.
-          The data type is uint32 and the shape is :math:`(n, 3)`.
-        - **scaler_f** (Tensor) - The 3-D scale factor (x, y, z),
-          between the real space float coordinates and the unsigned int coordinates.
-          The data type is float32 and the shape is :math:`(3,)`.
-        - **atom_a** (Tensor) - The first atom index of each bond.
-          The data type is int32 and the shape is :math:`(m,)`.
-        - **atom_b** (Tensor) - The second atom index of each bond.
-          The data type is int32 and the shape is :math:`(m,)`.
-        - **bond_k** (Tensor) - The force constant of each bond.
-          The data type is float32 and the shape is :math:`(m,)`.
-        - **bond_r0** (Tensor) - The equlibrium length of each bond.
-          The data type is float32 and the shape is :math:`(m,)`.
-
-    Outputs:
-        - **frc_f** (Tensor) - Same as operator BondForce().
-          The data type is float32 and the shape is :math:`(n, 3)`.
-        - **atom_v** (Tensor) - The accumulated virial coefficient for each atom.
-          The data type is float32 and the shape is :math:`(n,)`.
-
-    Supported Platforms:
-        ``GPU``
-    """
-
-    @prim_attr_register
-    def __init__(self, bond_numbers, atom_numbers):
-        """Initialize BondForceWithAtomVirial."""
-        validator.check_value_type('bond_numbers', bond_numbers, int, self.name)
-        validator.check_value_type('atom_numbers', atom_numbers, int, self.name)
-        self.bond_numbers = bond_numbers
-        self.atom_numbers = atom_numbers
-        self.add_prim_attr('bond_numbers', self.bond_numbers)
-        self.add_prim_attr('atom_numbers', self.atom_numbers)
-        self.init_prim_io_names(inputs=['uint_crd_f', 'scaler_f', 'atom_a', 'atom_b', 'bond_k', 'bond_r0'],
-                                outputs=['frc_f', 'atom_v'])
-
-    def infer_shape(self, uint_crd_f_shape, scaler_f_shape, atom_a_shape, atom_b_shape, bond_k_shape, bond_r0_shape):
-        cls_name = self.name
-        n = self.atom_numbers
-        m = self.bond_numbers
-        validator.check_int(len(uint_crd_f_shape), 2, Rel.EQ, "uint_crd_f_dim", cls_name)
-        validator.check_int(len(scaler_f_shape), 1, Rel.EQ, "scaler_f_dim", cls_name)
-        validator.check_int(len(atom_a_shape), 1, Rel.EQ, "atom_a_dim", cls_name)
-        validator.check_int(len(atom_b_shape), 1, Rel.EQ, "atom_b_dim", cls_name)
-        validator.check_int(len(bond_k_shape), 1, Rel.EQ, "bond_k_dim", cls_name)
-        validator.check_int(len(bond_r0_shape), 1, Rel.EQ, "bond_r0_dim", cls_name)
-        validator.check_int(uint_crd_f_shape[0], n, Rel.EQ, "uint_crd_f_shape[0]", cls_name)
-        validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape[1]", cls_name)
-        validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f_shape", cls_name)
-        validator.check_int(atom_a_shape[0], m, Rel.EQ, "uint_crd_f_shape", cls_name)
-        validator.check_int(atom_b_shape[0], m, Rel.EQ, "atom_b_shape", cls_name)
-        validator.check_int(bond_k_shape[0], m, Rel.EQ, "bond_k_shape", cls_name)
-        validator.check_int(bond_r0_shape[0], m, Rel.EQ, "bond_r0_shape", cls_name)
-
-        return uint_crd_f_shape, [n,]
-
-    def infer_dtype(self, uint_crd_f_dtype, scaler_f_type, atom_a_type, atom_b_type, bond_k_type, bond_r0_type):
-        validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name)
-        validator.check_tensor_dtype_valid('scaler_f', scaler_f_type, [mstype.float32], self.name)
-
-        validator.check_tensor_dtype_valid('atom_a', atom_a_type, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('atom_b', atom_b_type, [mstype.int32], self.name)
-
-        validator.check_tensor_dtype_valid('bond_k', bond_k_type, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('bond_r0', bond_r0_type, [mstype.float32], self.name)
-        return bond_r0_type, bond_r0_type
-
-
-class DihedralForce(PrimitiveWithInfer):
-    """
-    Calculate the force exerted by the dihedral term which made of 4-atoms
-    on the corresponding atoms. Assume the number of dihedral terms is m and
-    the number of atoms is n.
-
-    Because there is a large amount of inputs and each of them are related,
-    there is no way to construct `Examples` using random methods. For details, refer the webpage `SPONGE in MindSpore
-    <https://gitee.com/mindspore/mindscience/blob/master/MindSPONGE/docs/simple_formula.md>`_.
-
-    Args:
-        dihedral_numbers(int32): the number of dihedral terms m.
-
-    Inputs:
-        - **uint_crd_f** (Tensor) - The unsigned int coordinates
-          value of each atom. The data type is uint32 and the shape is :math:`(n, 3)`.
-        - **scaler_f** (Tensor) - The 3-D scale factor between
-          the real space float coordinates and the unsigned int coordinates.
-          The data type is float32 and the shape is :math:`(3,)`.
-        - **atom_a** (Tensor) - The 1st atom index of each dihedral.
-          The data type is int32 and the shape is :math:`(m,)`.
-        - **atom_b** (Tensor) - The 2nd atom index of each dihedral.
-          The data type is int32 and the shape is :math:`(m,)`.
-        - **atom_c** (Tensor) - The 3rd atom index of each dihedral.
-          The data type is int32 and the shape is :math:`(m,)`.
-        - **atom_d** (Tensor) - The 4th atom index of each dihedral.
-          4 atoms are connected in the form a-b-c-d.
-          The data type is int32 and the shape is :math:`(m,)`.
-        - **ipn** (Tensor) - The period of dihedral angle of each dihedral.
-          The data type is int32 and the shape is :math:`(m,)`.
-        - **pk** (Tensor) - The force constant of each dihedral.
-          The data type is float32 and the shape is :math:`(m,)`.
-        - **gamc** (Tensor) - k*cos(phi_0) of each dihedral.
-          The data type is float32 and the shape is :math:`(m,)`.
-        - **gams** (Tensor) - k*sin(phi_0) of each dihedral.
-          The data type is float32 and the shape is :math:`(m,)`.
-        - **pn** (Tensor) - The floating point form of ipn.
-          The data type is float32 and the shape is :math:`(m,)`.
-
-    Outputs:
-        - **frc_f** (Tensor) - The force felt by each atom.
-          The data type is float32 and the shape is :math:`(n, 3)`.
-
-    Supported Platforms:
-        ``GPU``
-    """
-
-    @prim_attr_register
-    def __init__(self, dihedral_numbers):
-        """Initialize DihedralForce."""
-        validator.check_value_type('dihedral_numbers', dihedral_numbers, int, self.name)
-        self.dihedral_numbers = dihedral_numbers
-        self.init_prim_io_names(inputs=['uint_crd_f', 'scaler_f', 'atom_a', 'atom_b', 'atom_c', 'atom_d', 'ipn', 'pk',
-                                        'gamc', 'gams', 'pn'],
-                                outputs=['frc_f'])
-        self.add_prim_attr('dihedral_numbers', self.dihedral_numbers)
-
-    def infer_shape(self, uint_crd_f_shape, scaler_f_shape, atom_a_shape, atom_b_shape, atom_c_shape, atom_d_shape,
-                    ipn_shape, pk_shape, gamc_shape, gams_shape, pn_shape):
-        cls_name = self.name
-        m = self.dihedral_numbers
-        validator.check_int(len(uint_crd_f_shape), 2, Rel.EQ, "uint_crd_f_dim", cls_name)
-        validator.check_int(len(scaler_f_shape), 1, Rel.EQ, "scaler_f_dim", cls_name)
-        validator.check_int(len(atom_a_shape), 1, Rel.EQ, "atom_a_dim", cls_name)
-        validator.check_int(len(atom_b_shape), 1, Rel.EQ, "atom_b_dim", cls_name)
-        validator.check_int(len(atom_c_shape), 1, Rel.EQ, "atom_c_dim", cls_name)
-        validator.check_int(len(atom_d_shape), 1, Rel.EQ, "atom_d_dim", cls_name)
-        validator.check_int(len(ipn_shape), 1, Rel.EQ, "ipn_dim", cls_name)
-        validator.check_int(len(pk_shape), 1, Rel.EQ, "pk_dim", cls_name)
-        validator.check_int(len(gamc_shape), 1, Rel.EQ, "gamc_dim", cls_name)
-        validator.check_int(len(gams_shape), 1, Rel.EQ, "gams_dim", cls_name)
-        validator.check_int(len(pn_shape), 1, Rel.EQ, "pn_dim", cls_name)
-
-        validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape[1]", cls_name)
-        validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f_shape", cls_name)
-        validator.check_int(atom_a_shape[0], m, Rel.EQ, "atom_a_shape", cls_name)
-        validator.check_int(atom_b_shape[0], m, Rel.EQ, "atom_b_shape", cls_name)
-        validator.check_int(atom_c_shape[0], m, Rel.EQ, "atom_c_shape", cls_name)
-        validator.check_int(atom_d_shape[0], m, Rel.EQ, "atom_d_shape", cls_name)
-        validator.check_int(ipn_shape[0], m, Rel.EQ, "ipn_shape", cls_name)
-        validator.check_int(pk_shape[0], m, Rel.EQ, "pk_shape", cls_name)
-        validator.check_int(gamc_shape[0], m, Rel.EQ, "gamc_shape", cls_name)
-        validator.check_int(gams_shape[0], m, Rel.EQ, "gams_shape", cls_name)
-        validator.check_int(pn_shape[0], m, Rel.EQ, "pn_shape", cls_name)
-        return uint_crd_f_shape
-
-    def infer_dtype(self, uint_crd_f_dtype, scaler_f_type, atom_a_type, atom_b_type, atom_c_type, atom_d_type,
-                    ipn_type, pk_type, gamc_type, gams_type, pn_type):
-        validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name)
-        validator.check_tensor_dtype_valid('scaler_f', scaler_f_type, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('atom_a', atom_a_type, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('atom_b', atom_b_type, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('atom_c', atom_c_type, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('atom_d', atom_d_type, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('ipn', ipn_type, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('pk', pk_type, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('gamc', gamc_type, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('gams', gams_type, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('pn', pn_type, [mstype.float32], self.name)
-
-        return pn_type
-
-
-class DihedralEnergy(PrimitiveWithInfer):
-    """
-    Calculate the potential energy caused by dihedral terms for each 4-atom pair.
-    Assume our system has n atoms and m dihedral terms.
-
-    Because there is a large amount of inputs and each of them are related,
-    there is no way to construct `Examples` using random methods. For details, refer the webpage `SPONGE in MindSpore
-    <https://gitee.com/mindspore/mindscience/blob/master/MindSPONGE/docs/simple_formula.md>`_.
-
-    Args:
-        dihedral_numbers(int32): the number of dihedral terms m.
-
-    Inputs:
-        - **uint_crd_f** (Tensor) - The unsigned int coordinates
-          value of each atom.
-          The data type is uint32 and the shape is :math:`(n, 3)`.
-        - **scaler_f** (Tensor) - The 3-D scale factor between
-          the real space float coordinates and the unsigned int coordinates.
-          The data type is float32 and the shape is :math:`(3,)`.
-        - **atom_a** (Tensor) - The 1st atom index of each dihedral.
-          The data type is int32 and the shape is :math:`(m,)`.
-        - **atom_b** (Tensor) - The 2nd atom index of each dihedral.
-          The data type is int32 and the shape is :math:`(m,)`.
-        - **atom_c** (Tensor) - The 3rd atom index of each dihedral.
-          The data type is int32 and the shape is :math:`(m,)`.
-        - **atom_d** (Tensor) - The 4th atom index of each dihedral.
-          4 atoms are connected in the form a-b-c-d.
-          The data type is int32 and the shape is :math:`(m,)`.
-        - **ipn** (Tensor) - The period of dihedral angle of each dihedral.
-          The data type is int32 and the shape is :math:`(m,)`.
-        - **pk** (Tensor) - The force constant of each dihedral.
-          The data type is int32 and the shape is :math:`(m,)`.
-        - **gamc** (Tensor) - k*cos(phi_0) of each dihedral.
-          The data type is float32 and the shape is :math:`(m,)`.
-        - **gams** (Tensor) - k*sin(phi_0) of each dihedral.
-          The data type is float32 and the shape is :math:`(m,)`.
-        - **pn** (Tensor) - The floating point form of ipn.
-          The data type is float32 and the shape is :math:`(m,)`.
-
-    Outputs:
-        - **ene** (Tensor) - The potential energy for each
-          dihedral term. The data type is float32 and the shape is :math:`(m,)`.
-
-    Supported Platforms:
-        ``GPU``
-    """
-
-    @prim_attr_register
-    def __init__(self, dihedral_numbers):
-        """Initialize DihedralEnergy."""
-        validator.check_value_type('dihedral_numbers', dihedral_numbers, int, self.name)
-        self.dihedral_numbers = dihedral_numbers
-        self.init_prim_io_names(inputs=['uint_crd_f', 'scaler_f', 'atom_a', 'atom_b', 'atom_c', 'atom_d', 'ipn', 'pk',
-                                        'gamc', 'gams', 'pn'],
-                                outputs=['ene'])
-        self.add_prim_attr('dihedral_numbers', self.dihedral_numbers)
-
-    def infer_shape(self, uint_crd_f_shape, scaler_f_shape, atom_a_shape, atom_b_shape, atom_c_shape, atom_d_shape,
-                    ipn_shape, pk_shape, gamc_shape, gams_shape, pn_shape):
-        cls_name = self.name
-        m = self.dihedral_numbers
-        validator.check_int(len(uint_crd_f_shape), 2, Rel.EQ, "uint_crd_f_dim", cls_name)
-        validator.check_int(len(scaler_f_shape), 1, Rel.EQ, "scaler_f_dim", cls_name)
-        validator.check_int(len(atom_a_shape), 1, Rel.EQ, "atom_a_dim", cls_name)
-        validator.check_int(len(atom_b_shape), 1, Rel.EQ, "atom_b_dim", cls_name)
-        validator.check_int(len(atom_c_shape), 1, Rel.EQ, "atom_c_dim", cls_name)
-        validator.check_int(len(atom_d_shape), 1, Rel.EQ, "atom_d_dim", cls_name)
-        validator.check_int(len(ipn_shape), 1, Rel.EQ, "ipn_dim", cls_name)
-        validator.check_int(len(pk_shape), 1, Rel.EQ, "pk_dim", cls_name)
-        validator.check_int(len(gamc_shape), 1, Rel.EQ, "gamc_dim", cls_name)
-        validator.check_int(len(gams_shape), 1, Rel.EQ, "gams_dim", cls_name)
-        validator.check_int(len(pn_shape), 1, Rel.EQ, "pn_dim", cls_name)
-
-        validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape[1]", cls_name)
-        validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f_shape", cls_name)
-        validator.check_int(atom_a_shape[0], m, Rel.EQ, "atom_a_shape", cls_name)
-        validator.check_int(atom_b_shape[0], m, Rel.EQ, "atom_b_shape", cls_name)
-        validator.check_int(atom_c_shape[0], m, Rel.EQ, "atom_c_shape", cls_name)
-        validator.check_int(atom_d_shape[0], m, Rel.EQ, "atom_d_shape", cls_name)
-        validator.check_int(ipn_shape[0], m, Rel.EQ, "ipn_shape", cls_name)
-        validator.check_int(pk_shape[0], m, Rel.EQ, "pk_shape", cls_name)
-        validator.check_int(gamc_shape[0], m, Rel.EQ, "gamc_shape", cls_name)
-        validator.check_int(gams_shape[0], m, Rel.EQ, "gams_shape", cls_name)
-        validator.check_int(pn_shape[0], m, Rel.EQ, "pn_shape", cls_name)
-        return [m,]
-
-    def infer_dtype(self, uint_crd_f_dtype, scaler_f_type, atom_a_type, atom_b_type, atom_c_type, atom_d_type,
-                    ipn_type, pk_type, gamc_type, gams_type, pn_type):
-        validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name)
-        validator.check_tensor_dtype_valid('scaler_f', scaler_f_type, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('atom_a', atom_a_type, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('atom_b', atom_b_type, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('atom_c', atom_c_type, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('atom_d', atom_d_type, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('ipn', ipn_type, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('pk', pk_type, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('gamc', gamc_type, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('gams', gams_type, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('pn', pn_type, [mstype.float32], self.name)
-
-        return pn_type
-
-
-class DihedralAtomEnergy(PrimitiveWithInfer):
-    """
-    Add the potential energy caused by dihedral terms to the total potential
-    energy of each atom.
-
-    Because there is a large amount of inputs and each of them are related,
-    there is no way to construct `Examples` using random methods. For details, refer the webpage `SPONGE in MindSpore
-    <https://gitee.com/mindspore/mindscience/blob/master/MindSPONGE/docs/simple_formula.md>`_.
-
-    The calculation formula is the same as operator DihedralEnergy().
-
-    Args:
-        dihedral_numbers(int32): the number of dihedral terms m.
-
-    Inputs:
-        - **uint_crd_f** (Tensor) - The unsigned int coordinates
-          value of each atom. The data type is uint32 and the shape is :math:`(n, 3)`.
-        - **scaler_f** (Tensor) - The 3-D scale factor between
-          the real space float coordinates and the unsigned int coordinates.
-          The data type is float32 and the shape is :math:`(3,)`.
-        - **atom_a** (Tensor) - The 1st atom index of each dihedral.
-          The data type is int32 and the shape is :math:`(m,)`.
-        - **atom_b** (Tensor) - The 2nd atom index of each dihedral.
-          The data type is int32 and the shape is :math:`(m,)`.
-        - **atom_c** (Tensor) - The 3rd atom index of each dihedral.
-          The data type is int32 and the shape is :math:`(m,)`.
-        - **atom_d** (Tensor) - The 4th atom index of each dihedral.
-          4 atoms are connected in the form a-b-c-d. The data type is int32 and the shape is :math:`(m,)`.
-        - **ipn** (Tensor) - The period of dihedral angle of each dihedral.
-          The data type is int32 and the shape is :math:`(m,)`.
-        - **pk** (Tensor) - The force constant of each dihedral.
-          The data type is float32 and the shape is :math:`(m,)`.
-        - **gamc** (Tensor) - k*cos(phi_0) of each dihedral.
-          The data type is float32 and the shape is :math:`(m,)`.
-        - **gams** (Tensor) - k*sin(phi_0) of each dihedral.
-          The data type is float32 and the shape is :math:`(m,)`.
-        - **pn** (Tensor) - The floating point form of ipn.
-          The data type is float32 and the shape is :math:`(m,)`.
-
-    Outputs:
-        - **ene** (Tensor) - The accumulated potential
-          energy for each atom. The data type is float32 and the shape is :math:`(n,)`.
-
-    Supported Platforms:
-        ``GPU``
-    """
-
-    @prim_attr_register
-    def __init__(self, dihedral_numbers):
-        """Initialize DihedralAtomEnergy."""
-        validator.check_value_type('dihedral_numbers', dihedral_numbers, int, self.name)
-        self.dihedral_numbers = dihedral_numbers
-        self.init_prim_io_names(inputs=['uint_crd_f', 'scaler_f', 'atom_a', 'atom_b', 'atom_c', 'atom_d', 'ipn', 'pk',
-                                        'gamc', 'gams', 'pn'],
-                                outputs=['ene'])
-        self.add_prim_attr('dihedral_numbers', self.dihedral_numbers)
-
-    def infer_shape(self, uint_crd_f_shape, scaler_f_shape, atom_a_shape, atom_b_shape, atom_c_shape, atom_d_shape,
-                    ipn_shape, pk_shape, gamc_shape, gams_shape, pn_shape):
-        cls_name = self.name
-        n = uint_crd_f_shape[0]
-        m = self.dihedral_numbers
-        validator.check_int(len(uint_crd_f_shape), 2, Rel.EQ, "uint_crd_f_dim", cls_name)
-        validator.check_int(len(scaler_f_shape), 1, Rel.EQ, "scaler_f_dim", cls_name)
-        validator.check_int(len(atom_a_shape), 1, Rel.EQ, "atom_a_dim", cls_name)
-        validator.check_int(len(atom_b_shape), 1, Rel.EQ, "atom_b_dim", cls_name)
-        validator.check_int(len(atom_c_shape), 1, Rel.EQ, "atom_c_dim", cls_name)
-        validator.check_int(len(atom_d_shape), 1, Rel.EQ, "atom_d_dim", cls_name)
-        validator.check_int(len(ipn_shape), 1, Rel.EQ, "ipn_dim", cls_name)
-        validator.check_int(len(pk_shape), 1, Rel.EQ, "pk_dim", cls_name)
-        validator.check_int(len(gamc_shape), 1, Rel.EQ, "gamc_dim", cls_name)
-        validator.check_int(len(gams_shape), 1, Rel.EQ, "gams_dim", cls_name)
-        validator.check_int(len(pn_shape), 1, Rel.EQ, "pn_dim", cls_name)
-
-        validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape[1]", cls_name)
-        validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f_shape", cls_name)
-        validator.check_int(atom_a_shape[0], m, Rel.EQ, "atom_a_shape", cls_name)
-        validator.check_int(atom_b_shape[0], m, Rel.EQ, "atom_b_shape", cls_name)
-        validator.check_int(atom_c_shape[0], m, Rel.EQ, "atom_c_shape", cls_name)
-        validator.check_int(atom_d_shape[0], m, Rel.EQ, "atom_d_shape", cls_name)
-        validator.check_int(ipn_shape[0], m, Rel.EQ, "ipn_shape", cls_name)
-        validator.check_int(pk_shape[0], m, Rel.EQ, "pk_shape", cls_name)
-        validator.check_int(gamc_shape[0], m, Rel.EQ, "gamc_shape", cls_name)
-        validator.check_int(gams_shape[0], m, Rel.EQ, "gams_shape", cls_name)
-        validator.check_int(pn_shape[0], m, Rel.EQ, "pn_shape", cls_name)
-        return [n,]
-
-    def infer_dtype(self, uint_crd_f_dtype, scaler_f_type, atom_a_type, atom_b_type, atom_c_type, atom_d_type,
-                    ipn_type, pk_type, gamc_type, gams_type, pn_type):
-        validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name)
-        validator.check_tensor_dtype_valid('scaler_f', scaler_f_type, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('atom_a', atom_a_type, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('atom_b', atom_b_type, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('atom_c', atom_c_type, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('atom_d', atom_d_type, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('ipn', ipn_type, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('pk', pk_type, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('gamc', gamc_type, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('gams', gams_type, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('pn', pn_type, [mstype.float32], self.name)
-
-        return pn_type
-
-
-class DihedralForceWithAtomEnergy(PrimitiveWithInfer):
-    """
-    Calculate dihedral force and potential energy together.
-
-    The calculation formula is the same as operator DihedralForce() and DihedralEnergy().
-
-    Because there is a large amount of inputs and each of them are related,
-    there is no way to construct `Examples` using random methods. For details, refer the webpage `SPONGE in MindSpore
-    <https://gitee.com/mindspore/mindscience/blob/master/MindSPONGE/docs/simple_formula.md>`_.
-
-    Args:
-        dihedral_numbers(int32): the number of dihedral terms m.
-
-    Inputs:
-        - **uint_crd_f** (Tensor) - The unsigned int coordinates
-          value of each atom. The data type is uint32 and the shape is :math:`(n, 3)`.
-        - **scaler_f** (Tensor) - The 3-D scale factor between
-          the real space float coordinates and the unsigned int coordinates.
-          The data type is float32 and the shape is :math:`(3,)`.
-        - **atom_a** (Tensor) - The 1st atom index of each dihedral.
-          The data type is int32 and the shape is :math:`(m,)`.
-        - **atom_b** (Tensor) - The 2nd atom index of each dihedral.
-          The data type is int32 and the shape is :math:`(m,)`.
-        - **atom_c** (Tensor) - The 3rd atom index of each dihedral.
-          The data type is int32 and the shape is :math:`(m,)`.
-        - **atom_d** (Tensor) - The 4th atom index of each dihedral.
-          4 atoms are connected in the form a-b-c-d. The data type is int32 and the shape is :math:`(m,)`.
-        - **ipn** (Tensor) - The period of dihedral angle of each dihedral.
-          The data type is int32 and the shape is :math:`(m,)`.
-        - **pk** (Tensor) - The force constant of each dihedral.
-          The data type is float32 and the shape is :math:`(m,)`.
-        - **gamc** (Tensor) - k*cos(phi_0) of each dihedral.
-          The data type is float32 and the shape is :math:`(m,)`.
-        - **gams** (Tensor) - k*sin(phi_0) of each dihedral.
-          The data type is float32 and the shape is :math:`(m,)`.
-        - **pn** (Tensor) - The floating point form of ipn.
-          The data type is float32 and the shape is :math:`(m,)`.
-
-    Outputs:
-        - **frc_f** (Tensor) - Same as operator DihedralForce().
-          The data type is float32 and the shape is :math:`(n, 3)`.
-        - **ene** (Tensor) - Same as operator DihedralAtomEnergy().
-          The data type is float32 and the shape is :math:`(n,)`.
-
-    Supported Platforms:
-        ``GPU``
-    """
-
-    @prim_attr_register
-    def __init__(self, dihedral_numbers):
-        """Initialize DihedralForceWithAtomEnergy."""
-        validator.check_value_type('dihedral_numbers', dihedral_numbers, int, self.name)
-        self.dihedral_numbers = dihedral_numbers
-        self.init_prim_io_names(inputs=['uint_crd_f', 'scaler_f', 'atom_a', 'atom_b', 'atom_c', 'atom_d', 'ipn', 'pk',
-                                        'gamc', 'gams', 'pn'],
-                                outputs=['frc_f', 'ene'])
-        self.add_prim_attr('dihedral_numbers', self.dihedral_numbers)
-
-    def infer_shape(self, uint_crd_f_shape, scaler_f_shape, atom_a_shape, atom_b_shape, atom_c_shape, atom_d_shape,
-                    ipn_shape, pk_shape, gamc_shape, gams_shape, pn_shape):
-        cls_name = self.name
-        n = uint_crd_f_shape[0]
-        m = self.dihedral_numbers
-        validator.check_int(len(uint_crd_f_shape), 2, Rel.EQ, "uint_crd_f_dim", cls_name)
-        validator.check_int(len(scaler_f_shape), 1, Rel.EQ, "scaler_f_dim", cls_name)
-        validator.check_int(len(atom_a_shape), 1, Rel.EQ, "atom_a_dim", cls_name)
-        validator.check_int(len(atom_b_shape), 1, Rel.EQ, "atom_b_dim", cls_name)
-        validator.check_int(len(atom_c_shape), 1, Rel.EQ, "atom_c_dim", cls_name)
-        validator.check_int(len(atom_d_shape), 1, Rel.EQ, "atom_d_dim", cls_name)
-        validator.check_int(len(ipn_shape), 1, Rel.EQ, "ipn_dim", cls_name)
-        validator.check_int(len(pk_shape), 1, Rel.EQ, "pk_dim", cls_name)
-        validator.check_int(len(gamc_shape), 1, Rel.EQ, "gamc_dim", cls_name)
-        validator.check_int(len(gams_shape), 1, Rel.EQ, "gams_dim", cls_name)
-        validator.check_int(len(pn_shape), 1, Rel.EQ, "pn_dim", cls_name)
-
-        validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape[1]", cls_name)
-        validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f_shape", cls_name)
-        validator.check_int(atom_a_shape[0], m, Rel.EQ, "atom_a_shape", cls_name)
-        validator.check_int(atom_b_shape[0], m, Rel.EQ, "atom_b_shape", cls_name)
-        validator.check_int(atom_c_shape[0], m, Rel.EQ, "atom_c_shape", cls_name)
-        validator.check_int(atom_d_shape[0], m, Rel.EQ, "atom_d_shape", cls_name)
-        validator.check_int(ipn_shape[0], m, Rel.EQ, "ipn_shape", cls_name)
-        validator.check_int(pk_shape[0], m, Rel.EQ, "pk_shape", cls_name)
-        validator.check_int(gamc_shape[0], m, Rel.EQ, "gamc_shape", cls_name)
-        validator.check_int(gams_shape[0], m, Rel.EQ, "gams_shape", cls_name)
-        validator.check_int(pn_shape[0], m, Rel.EQ, "pn_shape", cls_name)
-        return uint_crd_f_shape, [n,]
-
-    def infer_dtype(self, uint_crd_f_dtype, scaler_f_type, atom_a_type, atom_b_type, atom_c_type, atom_d_type,
-                    ipn_type, pk_type, gamc_type, gams_type, pn_type):
-        validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name)
-        validator.check_tensor_dtype_valid('scaler_f', scaler_f_type, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('atom_a', atom_a_type, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('atom_b', atom_b_type, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('atom_c', atom_c_type, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('atom_d', atom_d_type, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('ipn', ipn_type, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('pk', pk_type, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('gamc', gamc_type, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('gams', gams_type, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('pn', pn_type, [mstype.float32], self.name)
-
-        return pn_type, pn_type
-
-
-class AngleForce(PrimitiveWithInfer):
-    """
-    Calculate the force exerted by angles made of 3 atoms on the
-    corresponding atoms. Assume the number of angles is m and the
-    number of atoms is n.
-
-    Because there is a large amount of inputs and each of them are related,
-    there is no way to construct `Examples` using random methods. For details, refer the webpage `SPONGE in MindSpore
-    <https://gitee.com/mindspore/mindscience/blob/master/MindSPONGE/docs/simple_formula.md>`_.
-
-    .. math::
-        dr_{ab} = (x_b-x_a, y_b-y_a, z_b-z_a)
-    .. math::
-        dr_{cb} = (x_b-x_c, y_b-y_c, z_b-z_c)
-    .. math::
-        theta = arccos(inner_product(dr_{ab}, dr_{cb})/|dr_{ab}|/|dr_{cb}|)
-    .. math::
-        F_a = -2*k*(theta-theta_0)/sin(theta)*[cos(theta)/|dr_{ab}|^2*dr_{ab}
-            - 1/|dr_{ab}|/|dr_{cb}|*dr_{cb}]
-    .. math::
-        F_c = -2*k*(theta-theta_0)/sin(theta)*[cos(theta)/|dr_{cb}|^2*dr_{cb}
-            - 1/|dr_{cb}|/|dr_{ab}|*dr_{ab}]
-    .. math::
-        F_b = -F_a - F_c
-
-    Args:
-        angle_numbers(int32): the number of angles m.
-
-    Inputs:
-        - **uint_crd_f** (Tensor) - The unsigned int coordinate value of each atom.
-          The data type is uint32 and the shape is :math:`(n, 3)`.
-        - **scaler_f** (Tensor) - The 3-D scale factor between
-          the real space float coordinates and the unsigned int coordinates.
-          The data type is float32 and the shape is :math:`(3,)`.
-        - **atom_a** (Tensor) - The 1st atom index of each angle.
-          The data type is int32 and the shape is :math:`(m,)`.
-        - **atom_b** (Tensor) - The 2nd and the central atom index of each angle.
-          The data type is int32 and the shape is :math:`(m,)`.
-        - **atom_c** (Tensor) - The 3rd atom index of each angle.
-          The data type is int32 and the shape is :math:`(m,)`.
-        - **angle_k** (Tensor) - The force constant for each angle.
-          The data type is float32 and the shape is :math:`(m,)`.
-        - **angle_theta0** (Tensor) - The equilibrium position value for each angle.
-          The data type is float32 and the shape is :math:`(m,)`.
-
-    Outputs:
-        - **frc_f** (Tensor) - The force felt by each atom.
-          The data type is float32 and the shape is :math:`(n, 3)`.
-
-    Supported Platforms:
-        ``GPU``
-    """
-
-    @prim_attr_register
-    def __init__(self, angle_numbers):
-        """Initialize AngleForce."""
-        validator.check_value_type('angle_numbers', angle_numbers, int, self.name)
-        self.angle_numbers = angle_numbers
-        self.init_prim_io_names(inputs=['uint_crd_f', 'scaler_f', 'atom_a', 'atom_b', 'atom_c', 'angle_k',
-                                        'angle_theta0'],
-                                outputs=['frc_f'])
-        self.add_prim_attr('angle_numbers', self.angle_numbers)
-
-    def infer_shape(self, uint_crd_f_shape, scaler_f_shape, atom_a_shape, atom_b_shape, atom_c_shape, angle_k_shape,
-                    angle_theta0_shape):
-        cls_name = self.name
-        m = self.angle_numbers
-        validator.check_int(len(uint_crd_f_shape), 2, Rel.EQ, "uint_crd_f_dim", cls_name)
-        validator.check_int(len(scaler_f_shape), 1, Rel.EQ, "scaler_f_dim", cls_name)
-        validator.check_int(len(atom_a_shape), 1, Rel.EQ, "atom_a_dim", cls_name)
-        validator.check_int(len(atom_b_shape), 1, Rel.EQ, "atom_b_dim", cls_name)
-        validator.check_int(len(atom_c_shape), 1, Rel.EQ, "atom_c_dim", cls_name)
-        validator.check_int(len(angle_k_shape), 1, Rel.EQ, "angle_k_dim", cls_name)
-        validator.check_int(len(angle_theta0_shape), 1, Rel.EQ, "angle_theta0_dim", cls_name)
-
-        validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape[1]", cls_name)
-        validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f_shape", cls_name)
-        validator.check_int(atom_a_shape[0], m, Rel.EQ, "atom_a_shape", cls_name)
-        validator.check_int(atom_b_shape[0], m, Rel.EQ, "atom_b_shape", cls_name)
-        validator.check_int(atom_c_shape[0], m, Rel.EQ, "atom_c_shape", cls_name)
-        validator.check_int(angle_k_shape[0], m, Rel.EQ, "angle_k_shape", cls_name)
-        validator.check_int(angle_theta0_shape[0], m, Rel.EQ, "angle_theta0_shape", cls_name)
-        return uint_crd_f_shape
-
-    def infer_dtype(self, uint_crd_f_dtype, scaler_f_type, atom_a_type, atom_b_type, atom_c_type, angle_k_type,
-                    angle_theta0_type):
-        validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name)
-        validator.check_tensor_dtype_valid('scaler_f', scaler_f_type, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('atom_a', atom_a_type, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('atom_b', atom_b_type, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('atom_c', atom_c_type, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('angle_k', angle_k_type, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('angle_theta0', angle_theta0_type, [mstype.float32], self.name)
-        return angle_k_type
-
-
-class AngleEnergy(PrimitiveWithInfer):
-    """
-    Calculate the energy caused by 3-atoms angle term. Assume the number of angles is m and the
-    number of atoms is n.
-
-    Because there is a large amount of inputs and each of them are related,
-    there is no way to construct `Examples` using random methods. For details, refer the webpage `SPONGE in MindSpore
-    <https://gitee.com/mindspore/mindscience/blob/master/MindSPONGE/docs/simple_formula.md>`_.
-
-    .. math::
-        dr_{ab} = (x_b-x_a, y_b-y_a, z_b-z_a)
-    .. math::
-        dr_{cb} = (x_b-x_c, y_b-y_c, z_b-z_c)
-    .. math::
-        theta = arccos(inner_product(dr_{ab}, dr_{cb})/|dr_{ab}|/|dr_{cb}|)
-    .. math::
-        E = k*(theta - theta_0)^2
-
-    Args:
-        angle_numbers(int32): the number of angles m.
-
-    Inputs:
-        - **uint_crd_f** (Tensor) - The unsigned int coordinate value of each atom.
-          The data type is uint32 and the shape is :math:`(n, 3)`.
-        - **scaler_f** (Tensor) - The 3-D scale factor between
-          the real space float coordinates and the unsigned int coordinates.
-          The data type is float32 and the shape is :math:`(3,)`.
-        - **atom_a** (Tensor) - The 1st atom index of each angle.
-          The data type is int32 and the shape is :math:`(m,)`.
-        - **atom_b** (Tensor) - The 2nd and the central atom index of each angle.
-          The data type is int32 and the shape is :math:`(m,)`.
-        - **atom_c** (Tensor) - The 3rd atom index of each angle.
-          The data type is int32 and the shape is :math:`(m,)`.
-        - **angle_k** (Tensor) - The force constant for each angle.
-          The data type is float32 and the shape is :math:`(m,)`.
-        - **angle_theta0** (Tensor) - The equilibrium position value for each angle.
-          The data type is float32 and the shape is :math:`(m,)`.
-
-    Outputs:
-        - **ene** (Tensor) - The potential energy for each angle term.
-          The data type is float32 and the shape is :math:`(m,)`.
-
-    Supported Platforms:
-        ``GPU``
-    """
-
-    @prim_attr_register
-    def __init__(self, angle_numbers):
-        """Initialize AngleEnergy."""
-        validator.check_value_type('angle_numbers', angle_numbers, int, self.name)
-        self.angle_numbers = angle_numbers
-        self.init_prim_io_names(inputs=['uint_crd_f', 'scaler_f', 'atom_a', 'atom_b', 'atom_c', 'angle_k',
-                                        'angle_theta0'],
-                                outputs=['ene'])
-        self.add_prim_attr('angle_numbers', self.angle_numbers)
-
-    def infer_shape(self, uint_crd_f_shape, scaler_f_shape, atom_a_shape, atom_b_shape, atom_c_shape, angle_k_shape,
-                    angle_theta0_shape):
-        cls_name = self.name
-        m = self.angle_numbers
-        validator.check_int(len(uint_crd_f_shape), 2, Rel.EQ, "uint_crd_f_dim", cls_name)
-        validator.check_int(len(scaler_f_shape), 1, Rel.EQ, "scaler_f_dim", cls_name)
-        validator.check_int(len(atom_a_shape), 1, Rel.EQ, "atom_a_dim", cls_name)
-        validator.check_int(len(atom_b_shape), 1, Rel.EQ, "atom_b_dim", cls_name)
-        validator.check_int(len(atom_c_shape), 1, Rel.EQ, "atom_c_dim", cls_name)
-        validator.check_int(len(angle_k_shape), 1, Rel.EQ, "angle_k_dim", cls_name)
-        validator.check_int(len(angle_theta0_shape), 1, Rel.EQ, "angle_theta0_dim", cls_name)
-
-        validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape[1]", cls_name)
-        validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f_shape", cls_name)
-        validator.check_int(atom_a_shape[0], m, Rel.EQ, "atom_a_shape", cls_name)
-        validator.check_int(atom_b_shape[0], m, Rel.EQ, "atom_b_shape", cls_name)
-        validator.check_int(atom_c_shape[0], m, Rel.EQ, "atom_c_shape", cls_name)
-        validator.check_int(angle_k_shape[0], m, Rel.EQ, "angle_k_shape", cls_name)
-        validator.check_int(angle_theta0_shape[0], m, Rel.EQ, "angle_theta0_shape", cls_name)
-        return [m,]
-
-    def infer_dtype(self, uint_crd_f_dtype, scaler_f_type, atom_a_type, atom_b_type, atom_c_type, angle_k_type,
-                    angle_theta0_type):
-        validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name)
-        validator.check_tensor_dtype_valid('scaler_f', scaler_f_type, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('atom_a', atom_a_type, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('atom_b', atom_b_type, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('atom_c', atom_c_type, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('angle_k', angle_k_type, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('angle_theta0', angle_theta0_type, [mstype.float32], self.name)
-        return angle_k_type
-
-
-class AngleAtomEnergy(PrimitiveWithInfer):
-    """
-    Add the potential energy caused by angle terms to the total potential
-    energy of each atom. Assume the number of angles is m and the
-    number of atoms is n.
-
-    The calculation formula is the same as operator AngleEnergy().
-
-    Because there is a large amount of inputs and each of them are related,
-    there is no way to construct `Examples` using random methods. For details, refer the webpage `SPONGE in MindSpore
-    <https://gitee.com/mindspore/mindscience/blob/master/MindSPONGE/docs/simple_formula.md>`_.
-
-    Args:
-        angle_numbers(int32): the number of angles m.
-
-    Inputs:
-        - **uint_crd_f** (Tensor) - The unsigned int coordinate value of each atom.
-          The data type is uint32 and the shape is :math:`(n, 3)`.
-        - **scaler_f** (Tensor) - The 3-D scale factor between
-          the real space float coordinates and the unsigned int coordinates.
-          The data type is float32 and the shape is :math:`(3,)`.
-        - **atom_a** (Tensor) - The 1st atom index of each angle.
-          The data type is int32 and the shape is :math:`(m,)`.
-        - **atom_b** (Tensor) - The 2nd and the central atom index of each angle.
-          The data type is int32 and the shape is :math:`(m,)`.
-        - **atom_c** (Tensor) - The 3rd atom index of each angle.
-          The data type is int32 and the shape is :math:`(m,)`.
-        - **angle_k** (Tensor) - The force constant for each angle.
-          The data type is float32 and the shape is :math:`(m,)`.
-        - **angle_theta0** (Tensor) - The equilibrium position value for each angle.
-          The data type is float32 and the shape is :math:`(m,)`.
-
-    Outputs:
-        - **ene** (Tensor) - The accumulated potential energy for each atom.
-          The data type is float32 and the shape is :math:`(n,)`.
-
-    Supported Platforms:
-        ``GPU``
-    """
-
-    @prim_attr_register
-    def __init__(self, angle_numbers):
-        """Initialize AngleAtomEnergy."""
-        validator.check_value_type('angle_numbers', angle_numbers, int, self.name)
-        self.angle_numbers = angle_numbers
-        self.init_prim_io_names(inputs=['uint_crd_f', 'scaler_f', 'atom_a', 'atom_b', 'atom_c', 'angle_k',
-                                        'angle_theta0'],
-                                outputs=['ene'])
-        self.add_prim_attr('angle_numbers', self.angle_numbers)
-
-    def infer_shape(self, uint_crd_f_shape, scaler_f_shape, atom_a_shape, atom_b_shape, atom_c_shape, angle_k_shape,
-                    angle_theta0_shape):
-        cls_name = self.name
-        n = uint_crd_f_shape[0]
-        m = self.angle_numbers
-        validator.check_int(len(uint_crd_f_shape), 2, Rel.EQ, "uint_crd_f_dim", cls_name)
-        validator.check_int(len(scaler_f_shape), 1, Rel.EQ, "scaler_f_dim", cls_name)
-        validator.check_int(len(atom_a_shape), 1, Rel.EQ, "atom_a_dim", cls_name)
-        validator.check_int(len(atom_b_shape), 1, Rel.EQ, "atom_b_dim", cls_name)
-        validator.check_int(len(atom_c_shape), 1, Rel.EQ, "atom_c_dim", cls_name)
-        validator.check_int(len(angle_k_shape), 1, Rel.EQ, "angle_k_dim", cls_name)
-        validator.check_int(len(angle_theta0_shape), 1, Rel.EQ, "angle_theta0_dim", cls_name)
-
-        validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape[1]", cls_name)
-        validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f_shape", cls_name)
-        validator.check_int(atom_a_shape[0], m, Rel.EQ, "atom_a_shape", cls_name)
-        validator.check_int(atom_b_shape[0], m, Rel.EQ, "atom_b_shape", cls_name)
-        validator.check_int(atom_c_shape[0], m, Rel.EQ, "atom_c_shape", cls_name)
-        validator.check_int(angle_k_shape[0], m, Rel.EQ, "angle_k_shape", cls_name)
-        validator.check_int(angle_theta0_shape[0], m, Rel.EQ, "angle_theta0_shape", cls_name)
-        return [n,]
-
-    def infer_dtype(self, uint_crd_f_dtype, scaler_f_type, atom_a_type, atom_b_type, atom_c_type, angle_k_type,
-                    angle_theta0_type):
-        validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name)
-        validator.check_tensor_dtype_valid('scaler_f', scaler_f_type, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('atom_a', atom_a_type, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('atom_b', atom_b_type, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('atom_c', atom_c_type, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('angle_k', angle_k_type, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('angle_theta0', angle_theta0_type, [mstype.float32], self.name)
-        return angle_k_type
-
-
-class AngleForceWithAtomEnergy(PrimitiveWithInfer):
-    """
-    Calculate angle force and potential energy together. Assume the number of angles is m and the
-    number of atoms is n.
-
-    The calculation formula is the same as operator AngleForce() and AngleEnergy().
-
-    Because there is a large amount of inputs and each of them are related,
-    there is no way to construct `Examples` using random methods. For details, refer the webpage `SPONGE in MindSpore
-    <https://gitee.com/mindspore/mindscience/blob/master/MindSPONGE/docs/simple_formula.md>`_.
-
-    Args:
-        angle_numbers(int32): the number of angles m.
-
-    Inputs:
-        - **uint_crd_f** (Tensor) - The unsigned int coordinate value of each atom.
-          The data type is uint32 and the shape is :math:`(n, 3)`.
-        - **scaler_f** (Tensor) - The 3-D scale factor between
-          the real space float coordinates and the unsigned int coordinates.
-          The data type is float and the shape is :math:`(3,)`.
-        - **atom_a** (Tensor) - The 1st atom index of each angle.
-          The data type is int32 and the shape is :math:`(m,)`.
-        - **atom_b** (Tensor) - The 2nd and the central atom index of each angle.
-          The data type is int32 and the shape is :math:`(m,)`.
-        - **atom_c** (Tensor) - The 3rd atom index of each angle.
-          The data type is int32 and the shape is :math:`(m,)`.
-        - **angle_k** (Tensor) - The force constant for each angle.
-          The data type is float32 and the shape is :math:`(m,)`.
-        - **angle_theta0** (Tensor) - The equilibrium position value for each angle.
-          The data type is float32 and the shape is :math:`(m,)`.
-
-    Outputs:
-        - **frc_f** (Tensor) - same as operator AngleForce().
-          The data type is float32 and the shape is :math:`(n, 3)`.
-        - **ene** (Tensor) - same as operator AngleAtomEnergy().
-          The data type is float and the shape is :math:`(n,)`.
-
-    Supported Platforms:
-        ``GPU``
-    """
-
-    @prim_attr_register
-    def __init__(self, angle_numbers):
-        """Initialize AngleForceWithAtomEnergy."""
-        validator.check_value_type('angle_numbers', angle_numbers, int, self.name)
-        self.angle_numbers = angle_numbers
-        self.init_prim_io_names(inputs=['uint_crd_f', 'scaler_f', 'atom_a', 'atom_b', 'atom_c', 'angle_k',
-                                        'angle_theta0'],
-                                outputs=['frc_f', 'ene'])
-        self.add_prim_attr('angle_numbers', self.angle_numbers)
-
-    def infer_shape(self, uint_crd_f_shape, scaler_f_shape, atom_a_shape, atom_b_shape, atom_c_shape, angle_k_shape,
-                    angle_theta0_shape):
-        cls_name = self.name
-        n = uint_crd_f_shape[0]
-        m = self.angle_numbers
-        validator.check_int(len(uint_crd_f_shape), 2, Rel.EQ, "uint_crd_f_dim", cls_name)
-        validator.check_int(len(scaler_f_shape), 1, Rel.EQ, "scaler_f_dim", cls_name)
-        validator.check_int(len(atom_a_shape), 1, Rel.EQ, "atom_a_dim", cls_name)
-        validator.check_int(len(atom_b_shape), 1, Rel.EQ, "atom_b_dim", cls_name)
-        validator.check_int(len(atom_c_shape), 1, Rel.EQ, "atom_c_dim", cls_name)
-        validator.check_int(len(angle_k_shape), 1, Rel.EQ, "angle_k_dim", cls_name)
-        validator.check_int(len(angle_theta0_shape), 1, Rel.EQ, "angle_theta0_dim", cls_name)
-
-        validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape[1]", cls_name)
-        validator.check_int(scaler_f_shape[0], 3, Rel.EQ, "scaler_f_shape", cls_name)
-        validator.check_int(atom_a_shape[0], m, Rel.EQ, "atom_a_shape", cls_name)
-        validator.check_int(atom_b_shape[0], m, Rel.EQ, "atom_b_shape", cls_name)
-        validator.check_int(atom_c_shape[0], m, Rel.EQ, "atom_c_shape", cls_name)
-        validator.check_int(angle_k_shape[0], m, Rel.EQ, "angle_k_shape", cls_name)
-        validator.check_int(angle_theta0_shape[0], m, Rel.EQ, "angle_theta0_shape", cls_name)
-        return uint_crd_f_shape, [n,]
-
-    def infer_dtype(self, uint_crd_f_dtype, scaler_f_type, atom_a_type, atom_b_type, atom_c_type, angle_k_type,
-                    angle_theta0_type):
-        validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name)
-        validator.check_tensor_dtype_valid('scaler_f', scaler_f_type, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('atom_a', atom_a_type, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('atom_b', atom_b_type, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('atom_c', atom_c_type, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('angle_k', angle_k_type, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('angle_theta0', angle_theta0_type, [mstype.float32], self.name)
-        return angle_k_type, angle_k_type
-
-
-class Dihedral14LJForce(PrimitiveWithInfer):
-    """
-    Calculate the Lennard-Jones part of 1,4 dihedral force correction
-    for each necessary dihedral terms on the corresponding atoms.
-
-    Assume the number of necessary dihedral 1,4 terms is m, the number of atoms is n,
-    and the number of Lennard-Jones types for all atoms is P, which means
-    there will be q = P*(P+1)/2 types of possible Lennard-Jones interactions
-    for all kinds of atom pairs.
-
-    Because there is a large amount of inputs and each of them are related,
-    there is no way to construct `Examples` using random methods. For details, refer the webpage `SPONGE in MindSpore
-    <https://gitee.com/mindspore/mindscience/blob/master/MindSPONGE/docs/simple_formula.md>`_.
-
-    .. math::
-        dr = (x_a-x_b, y_a-y_b, z_a-z_b)
-    .. math::
-        F = k*(-12*A/|dr|^{14} + 6*B/|dr|^{8})*dr
-
-    Args:
-        nb14_numbers (int32): the number of necessary dihedral 1,4 terms m.
-        atom_numbers (int32): the number of atoms n.
-
-    Inputs:
-        - **uint_crd_f** (Tensor) - The unsigned int coordinate value of each atom.
-          The data type is uint32 and the shape is :math:`(n, 3)`.
-        - **LJ_type** (Tensor) - The Lennard-Jones type of each atom.
-          The data type is int32 and the shape is :math:`(n,)`.
-        - **charge** (Tensor) - The charge of each atom.
-          The data type is float32 and the shape is :math:`(n,)`.
-        - **boxlength_f** (Tensor) - The length of molecular simulation box in 3 dimensions.
-          The data type is float32 and the shape is :math:`(3,)`.
-        - **a_14** (Tensor) - The first atom index of each dihedral 1,4 term.
-          The data type is int32 and the shape is :math:`(m,)`.
-        - **b_14** (Tensor) - The second atom index of each dihedral 1,4 term.
-          The data type is int32 and the shape is :math:`(m,)`.
-        - **lj_scale_factor** (Tensor) - The scale factor for the
-          Lennard-Jones part of force correction of each dihedral 1,4 term.
-          The data type is float32 and the shape is :math:`(m,)`.
-        - **LJ_type_A** (Tensor) - The A parameter in Lennard-Jones scheme of each atom pair type.
-          q is the number of atom pair. The data type is float32 and the shape is :math:`(q,)`.
-        - **LJ_type_B** (Tensor) - The B parameter in Lennard-Jones scheme of each atom pair type.
-          q is the number of atom pair. The data type is float32 and the shape is :math:`(q,)`.
-
-    Outputs:
-        - **frc_f** (Tensor) - The force felt by each atom.
-          The data type is float32 and the shape is :math:`(n, 3)`.
-
-    Supported Platforms:
-        ``GPU``
-    """
-
-    @prim_attr_register
-    def __init__(self, nb14_numbers, atom_numbers):
-        """Initialize Dihedral14LJForce."""
-        validator.check_value_type('nb14_numbers', nb14_numbers, int, self.name)
-        validator.check_value_type('atom_numbers', atom_numbers, int, self.name)
-        self.dihedral_14_numbers = nb14_numbers
-        self.atom_numbers = atom_numbers
-        self.init_prim_io_names(
-            inputs=['uint_crd_f', 'LJtype', 'charge', 'boxlength_f', 'a_14', 'b_14', 'lj_scale_factor',
-                    'LJ_type_A', 'LJ_type_B'],
-            outputs=['frc_f'])
-        self.add_prim_attr('dihedral_14_numbers', self.dihedral_14_numbers)
-        self.add_prim_attr('atom_numbers', self.atom_numbers)
-
-    def infer_shape(self, uint_crd_f_shape, ljtype_shape, charge_shape, boxlength_f_shape, a_14_shape, b_14_shape,
-                    lj_scale_factor_shape, lj_type_a_shape, lj_type_b_shape):
-        cls_name = self.name
-        n = self.atom_numbers
-        m = self.dihedral_14_numbers
-        q = lj_type_a_shape[0]
-        validator.check_int(len(uint_crd_f_shape), 2, Rel.EQ, "uint_crd_f_dim", cls_name)
-        validator.check_int(len(ljtype_shape), 1, Rel.EQ, "LJtype_dim", cls_name)
-        validator.check_int(len(charge_shape), 1, Rel.EQ, "charge_dim", cls_name)
-        validator.check_int(len(boxlength_f_shape), 1, Rel.EQ, "boxlength_f_dim", cls_name)
-        validator.check_int(len(a_14_shape), 1, Rel.EQ, "a_14_dim", cls_name)
-        validator.check_int(len(b_14_shape), 1, Rel.EQ, "b_14_dim", cls_name)
-        validator.check_int(len(lj_scale_factor_shape), 1, Rel.EQ, "lj_scale_factor_dim", cls_name)
-        validator.check_int(len(lj_type_b_shape), 1, Rel.EQ, "LJ_type_B_dim", cls_name)
-
-        validator.check_int(uint_crd_f_shape[0], n, Rel.EQ, "uint_crd_f[0]", cls_name)
-        validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f[1]", cls_name)
-        validator.check_int(ljtype_shape[0], n, Rel.EQ, "LJtype", cls_name)
-        validator.check_int(charge_shape[0], n, Rel.EQ, "charge", cls_name)
-        validator.check_int(boxlength_f_shape[0], 3, Rel.EQ, "boxlength_f", cls_name)
-        validator.check_int(lj_type_b_shape[0], q, Rel.EQ, "LJ_type_B", cls_name)
-        validator.check_int(a_14_shape[0], m, Rel.EQ, "a_14_shape", cls_name)
-        validator.check_int(b_14_shape[0], m, Rel.EQ, "b_14_shape", cls_name)
-        validator.check_int(lj_scale_factor_shape[0], m, Rel.EQ, "lj_scale_factor_shape", cls_name)
-        return uint_crd_f_shape
-
-    def infer_dtype(self, uint_crd_f_dtype, ljtype_dtype, charge_dtype, boxlength_f_type, a_14_type, b_14_type,
-                    lj_scale_factor_type, lj_type_a_type, lj_type_b_type):
-        validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name)
-        validator.check_tensor_dtype_valid('LJtype', ljtype_dtype, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('charge', charge_dtype, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('boxlength_f', boxlength_f_type, [mstype.float32], self.name)
-
-        validator.check_tensor_dtype_valid('a_14', a_14_type, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('b_14', b_14_type, [mstype.int32], self.name)
-
-        validator.check_tensor_dtype_valid('lj_scale_factor', lj_scale_factor_type, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('LJ_type_A', lj_type_a_type, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('LJ_type_B', lj_type_b_type, [mstype.float32], self.name)
-        return lj_type_b_type
-
-
-class Dihedral14LJEnergy(PrimitiveWithInfer):
-    """
-    Calculate the Lennard-Jones part of 1,4 dihedral energy correction for
-    each necessary dihedral terms on the corresponding atoms.
-
-    Because there is a large amount of inputs and each of them are related,
-    there is no way to construct `Examples` using random methods. For details, refer the webpage `SPONGE in MindSpore
-    <https://gitee.com/mindspore/mindscience/blob/master/MindSPONGE/docs/simple_formula.md>`_.
-
-    .. math::
-        dr = (x_a-x_b, y_a-y_b, z_a-z-b)
-    .. math::
-        E = k*(A/|dr|^{12} - B/|dr|^{6})
-
-    Args:
-        nb14_numbers (int32): the number of necessary dihedral 1,4 terms m.
-        atom_numbers (int32): the number of atoms n.
-
-    Inputs:
-        - **uint_crd_f** (Tensor) - The unsigned int coordinate value of each atom.
-          The data type is uint32 and the shape is :math:`(n, 3)`.
-        - **LJ_type** (Tensor) - The Lennard-Jones type of each atom.
-          The data type is int32 and the shape is :math:`(n,)`.
-        - **charge** (Tensor) - The charge of each atom.
-          The data type is float32 and the shape is :math:`(n,)`.
-        - **boxlength_f** (Tensor) - The length of molecular simulation box in 3 dimensions.
-          The data type is float32 and the shape is :math:`(3,)`.
-        - **a_14** (Tensor) - The first atom index of each dihedral 1,4 term.
-          The data type is int32 and the shape is :math:`(m,)`.
-        - **b_14** (Tensor) - The second atom index of each dihedral 1,4 term.
-          The data type is int32 and the shape is :math:`(m,)`.
-        - **lj_scale_factor** (Tensor) - The scale factor for the
-          Lennard-Jones part of force correction of each dihedral 1,4 term.
-          The data type is float32 and the shape is :math:`(m,)`.
-        - **LJ_type_A** (Tensor) - The A parameter in Lennard-Jones scheme of each atom pair type.
-          q is the number of atom pair. The data type is float32 and the shape is :math:`(q,)`.
-        - **LJ_type_B** (Tensor) - The B parameter in Lennard-Jones scheme of each atom pair type.
-          q is the number of atom pair. The data type is float32 and the shape is :math:`(q,)`.
-
-    Outputs:
-        - **ene** (Tensor) - The Lennard-Jones potential energy correction.
-          The data type is float32 and the shape is :math:`(m,)`.
-
-    Supported Platforms:
-        ``GPU``
-    """
-
-    @prim_attr_register
-    def __init__(self, nb14_numbers, atom_numbers):
-        """Initialize Dihedral14LJEnergy"""
-        validator.check_value_type('nb14_numbers', nb14_numbers, int, self.name)
-        validator.check_value_type('atom_numbers', atom_numbers, int, self.name)
-        self.dihedral_14_numbers = nb14_numbers
-        self.atom_numbers = atom_numbers
-
-        self.init_prim_io_names(
-            inputs=['uint_crd_f', 'LJtype', 'charge', 'boxlength_f', 'a_14', 'b_14', 'lj_scale_factor',
-                    'LJ_type_A', 'LJ_type_B'],
-            outputs=['ene'])
-        self.add_prim_attr('dihedral_14_numbers', self.dihedral_14_numbers)
-        self.add_prim_attr('atom_numbers', self.atom_numbers)
-
-    def infer_shape(self, uint_crd_f_shape, ljtype_shape, charge_shape, boxlength_f_shape, a_14_shape, b_14_shape,
-                    lj_scale_factor_shape, lj_type_a_shape, lj_type_b_shape):
-        cls_name = self.name
-        n = self.atom_numbers
-        m = self.dihedral_14_numbers
-        q = lj_type_a_shape[0]
-        validator.check_int(len(uint_crd_f_shape), 2, Rel.EQ, "uint_crd_f_dim", cls_name)
-        validator.check_int(len(ljtype_shape), 1, Rel.EQ, "LJtype_dim", cls_name)
-        validator.check_int(len(charge_shape), 1, Rel.EQ, "charge_dim", cls_name)
-        validator.check_int(len(boxlength_f_shape), 1, Rel.EQ, "boxlength_f_dim", cls_name)
-        validator.check_int(len(a_14_shape), 1, Rel.EQ, "a_14_dim", cls_name)
-        validator.check_int(len(b_14_shape), 1, Rel.EQ, "b_14_dim", cls_name)
-        validator.check_int(len(lj_scale_factor_shape), 1, Rel.EQ, "lj_scale_factor_dim", cls_name)
-        validator.check_int(len(lj_type_b_shape), 1, Rel.EQ, "LJ_type_B_dim", cls_name)
-
-        validator.check_int(uint_crd_f_shape[0], n, Rel.EQ, "uint_crd_f[0]", cls_name)
-        validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f[1]", cls_name)
-        validator.check_int(ljtype_shape[0], n, Rel.EQ, "LJtype", cls_name)
-        validator.check_int(charge_shape[0], n, Rel.EQ, "charge", cls_name)
-        validator.check_int(boxlength_f_shape[0], 3, Rel.EQ, "boxlength_f", cls_name)
-        validator.check_int(lj_type_b_shape[0], q, Rel.EQ, "LJ_type_B", cls_name)
-        validator.check_int(a_14_shape[0], m, Rel.EQ, "a_14_shape", cls_name)
-        validator.check_int(b_14_shape[0], m, Rel.EQ, "b_14_shape", cls_name)
-        validator.check_int(lj_scale_factor_shape[0], m, Rel.EQ, "lj_scale_factor_shape", cls_name)
-        return [self.dihedral_14_numbers,]
-
-    def infer_dtype(self, uint_crd_f_dtype, ljtype_dtype, charge_dtype, boxlength_f_type, a_14_type, b_14_type,
-                    lj_scale_factor_type, lj_type_a_type, lj_type_b_type):
-        validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name)
-        validator.check_tensor_dtype_valid('LJtype', ljtype_dtype, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('charge', charge_dtype, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('boxlength_f', boxlength_f_type, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('a_14', a_14_type, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('b_14', b_14_type, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('lj_scale_factor', lj_scale_factor_type, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('LJ_type_A', lj_type_a_type, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('LJ_type_B', lj_type_b_type, [mstype.float32], self.name)
-
-        return lj_type_a_type
-
-
-class Dihedral14LJForceWithDirectCF(PrimitiveWithInfer):
-    """
-    Calculate the Lennard-Jones part and the Coulomb part of force correction
-    for each necessary dihedral 1,4 terms.
-
-    Because there is a large amount of inputs and each of them are related,
-    there is no way to construct `Examples` using random methods. For details, refer the webpage `SPONGE in MindSpore
-    <https://gitee.com/mindspore/mindscience/blob/master/MindSPONGE/docs/simple_formula.md>`_.
-
-    The calculation formula of the Lennard-Jones part is the same as operator
-    Dihedral14LJForce(), and the Coulomb part is as follows:
-
-    .. math::
-            dr = (x_a-x_b, y_a-y_b, z_a-z_b)
-    .. math::
-            F = -k*q_a*q_b/|r|^3*dr
-
-    Args:
-        nb14_numbers (int32): the number of necessary dihedral 1,4 terms m.
-        atom_numbers (int32): the number of atoms n.
-
-    Inputs:
-        - **uint_crd_f** (Tensor) - The unsigned int coordinate value of each atom.
-          The data type is uint32 and the shape is :math:`(n, 3)`.
-        - **LJ_type** (Tensor) - The Lennard-Jones type of each atom.
-          The data type is int32 and the shape is :math:`(n,)`.
-        - **charge** (Tensor) - The charge of each atom.
-          The data type is float32 and the shape is :math:`(n,)`.
-        - **boxlength_f** (Tensor) - The length of molecular simulation box in 3 dimensions.
-          The data type is float32 and the shape is :math:`(3,)`.
-        - **a_14** (Tensor) - The first atom index of each dihedral 1,4 term.
-          The data type is int32 and the shape is :math:`(m,)`.
-        - **b_14** (Tensor) - The second atom index of each dihedral 1,4 term.
-          The data type is int32 and the shape is :math:`(m,)`.
-        - **lj_scale_factor** (Tensor) - The scale factor for the
-          Lennard-Jones part of force correction of each dihedral 1,4 term.
-          The data type is float32 and the shape is :math:`(m,)`.
-        - **cf_scale_factor** (Tensor) - The scale factor for the
-          Coulomb part of force correction for each dihedral 1,4 terms.
-          The data type is float32 and the shape is :math:`(m,)`.
-        - **LJ_type_A** (Tensor) - The A parameter in Lennard-Jones scheme of each atom pair type.
-          q is the number of atom pair. The data type is float32 and the shape is :math:`(q,)`.
-        - **LJ_type_B** (Tensor) - The B parameter in Lennard-Jones shceme of each atom pair type.
-          q is the number of atom pair. The data type is float32 and the shape is :math:`(q,)`.
-
-    Outputs:
-        - **frc_f** (Tensor) - The force felt by each atom.
-          The data type is float32 and the shape is :math:`(n, 3)`
-
-    Supported Platforms:
-        ``GPU``
-    """
-
-    @prim_attr_register
-    def __init__(self, nb14_numbers, atom_numbers):
-        """Initialize Dihedral14LJForceWithDirectCF."""
-        validator.check_value_type('nb14_numbers', nb14_numbers, int, self.name)
-        validator.check_value_type('atom_numbers', atom_numbers, int, self.name)
-        self.dihedral_14_numbers = nb14_numbers
-        self.atom_numbers = atom_numbers
-
-        self.init_prim_io_names(
-            inputs=['uint_crd_f', 'LJtype', 'charge', 'boxlength_f', 'a_14', 'b_14', 'lj_scale_factor',
-                    'cf_scale_factor',
-                    'LJ_type_A', 'LJ_type_B'],
-            outputs=['frc_f'])
-        self.add_prim_attr('dihedral_14_numbers', self.dihedral_14_numbers)
-        self.add_prim_attr('atom_numbers', self.atom_numbers)
-
-    def infer_shape(self, uint_crd_f_shape, ljtype_shape, charge_shape, boxlength_f_shape, a_14_shape, b_14_shape,
-                    lj_scale_factor_shape, cf_scale_factor_shape, lj_type_a_shape, lj_type_b_shape):
-        cls_name = self.name
-        n = self.atom_numbers
-        m = self.dihedral_14_numbers
-        q = lj_type_a_shape[0]
-        validator.check_int(len(uint_crd_f_shape), 2, Rel.EQ, "uint_crd_f_dim", cls_name)
-        validator.check_int(len(ljtype_shape), 1, Rel.EQ, "LJtype_dim", cls_name)
-        validator.check_int(len(charge_shape), 1, Rel.EQ, "charge_dim", cls_name)
-        validator.check_int(len(boxlength_f_shape), 1, Rel.EQ, "boxlength_f_dim", cls_name)
-        validator.check_int(len(a_14_shape), 1, Rel.EQ, "a_14_dim", cls_name)
-        validator.check_int(len(b_14_shape), 1, Rel.EQ, "b_14_dim", cls_name)
-        validator.check_int(len(lj_scale_factor_shape), 1, Rel.EQ, "lj_scale_factor_dim", cls_name)
-        validator.check_int(len(cf_scale_factor_shape), 1, Rel.EQ, "cf_scale_factor_dim", cls_name)
-        validator.check_int(len(lj_type_b_shape), 1, Rel.EQ, "LJ_type_B_dim", cls_name)
-
-        validator.check_int(uint_crd_f_shape[0], n, Rel.EQ, "uint_crd_f_shape[0]", cls_name)
-        validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape[1]", cls_name)
-        validator.check_int(ljtype_shape[0], n, Rel.EQ, "LJtype_shape", cls_name)
-        validator.check_int(charge_shape[0], n, Rel.EQ, "charge_shape", cls_name)
-        validator.check_int(boxlength_f_shape[0], 3, Rel.EQ, "boxlength_f_shape", cls_name)
-        validator.check_int(lj_type_b_shape[0], q, Rel.EQ, "LJ_type_B_shape", cls_name)
-        validator.check_int(a_14_shape[0], m, Rel.EQ, "a_14_shape", cls_name)
-        validator.check_int(b_14_shape[0], m, Rel.EQ, "b_14_shape", cls_name)
-        validator.check_int(lj_scale_factor_shape[0], m, Rel.EQ, "lj_scale_factor_shape", cls_name)
-        validator.check_int(cf_scale_factor_shape[0], m, Rel.EQ, "cf_scale_factor_shape", cls_name)
-        return [self.atom_numbers, 3]
-
-    def infer_dtype(self, uint_crd_f_dtype, ljtype_dtype, charge_dtype, boxlength_f_type, a_14_type, b_14_type,
-                    lj_scale_factor_type, cf_scale_factor_type, lj_type_a_type, lj_type_b_type):
-        validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name)
-        validator.check_tensor_dtype_valid('LJtype', ljtype_dtype, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('charge', charge_dtype, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('boxlength_f', boxlength_f_type, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('a_14', a_14_type, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('b_14', b_14_type, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('lj_scale_factor', lj_scale_factor_type, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('cf_scale_factor', cf_scale_factor_type, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('LJ_type_A', lj_type_a_type, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('LJ_type_B', lj_type_b_type, [mstype.float32], self.name)
-
-        return lj_type_a_type
-
-
-class Dihedral14LJCFForceWithAtomEnergy(PrimitiveWithInfer):
-    """
-    Calculate the Lennard-Jones and Coulumb energy correction and force correction
-    for each necessary dihedral 1,4 terms together and add them to the total force
-    and potential energy for each atom.
-
-    The calculation formula of force correction is the same as operator
-    :class:`Dihedral14LJForceWithDirectCF`, and the energy correction part is the same
-    as operator :class:`Dihedral14LJEnergy` and :class:`Dihedral14CFEnergy`.
-
-    Because there is a large amount of inputs and each of them are related,
-    there is no way to construct `Examples` using random methods. For details, refer the webpage `SPONGE in MindSpore
-    <https://gitee.com/mindspore/mindscience/blob/master/MindSPONGE/docs/simple_formula.md>`_.
-
-    Args:
-        nb14_numbers (int32): the number of necessary dihedral 1,4 terms m.
-        atom_numbers (int32): the number of atoms n.
-
-    Inputs:
-        - **uint_crd_f** (Tensor) - The unsigned int coordinate value of each atom.
-          The data type is uint32 and the shape is :math:`(n, 3)`.
-        - **LJ_type** (Tensor) - The Lennard-Jones type of each atom.
-          The data type is int32 and the shape is :math:`(n,)`.
-        - **charge** (Tensor) - The charge of each atom.
-          The data type is float32 and the shape is :math:`(n,)`.
-        - **boxlength_f** (Tensor) - The length of molecular simulation box in 3 dimensions.
-          The data type is float32 and the shape is :math:`(3,)`.
-        - **a_14** (Tensor) - The first atom index of each dihedral 1,4 term.
-          The data type is int32 and the shape is :math:`(m,)`.
-        - **b_14** (Tensor) - The second atom index of each dihedral 1,4 term.
-          The data type is int32 and the shape is :math:`(m,)`.
-        - **lj_scale_factor** (Tensor) - The scale factor for the
-          Lennard-Jones part of force correction of each dihedral 1,4 term.
-          The data type is float32 and the shape is :math:`(m,)`.
-        - **cf_scale_factor** (Tensor) - The scale factor for the
-          Coulomb part of force correction for each dihedral 1,4 terms.
-          The data type is float32 and the shape is :math:`(m,)`.
-        - **LJ_type_A** (Tensor) - The A parameter in Lennard-Jones scheme of each atom pair type.
-          q is the number of atom pair. The data type is float32 and the shape is :math:`(q,)`.
-        - **LJ_type_B** (Tensor) - The B parameter in Lennard-Jones scheme of each atom pair type.
-          q is the number of atom pair. The data type is float32 and the shape is :math:`(q,)`.
-
-    Outputs:
-        - **frc_f** (Tensor) - The force felt by each atom.
-          The data type is float32 and the shape is :math:`(n, 3)`.
-        - **atom_energy** (Tensor) - The accumulated potential energy for each atom.
-          The data type is float32 and the shape is :math:`(n,)`.
-
-    Supported Platforms:
-        ``GPU``
-    """
-
-    @prim_attr_register
-    def __init__(self, nb14_numbers, atom_numbers):
-        """Initialize Dihedral14LJCFForceWithAtomEnergy."""
-        validator.check_value_type('nb14_numbers', nb14_numbers, int, self.name)
-        validator.check_value_type('atom_numbers', atom_numbers, int, self.name)
-        self.dihedral_14_numbers = nb14_numbers
-        self.atom_numbers = atom_numbers
-
-        self.init_prim_io_names(
-            inputs=['uint_crd_f', 'LJtype', 'charge', 'boxlength_f', 'a_14', 'b_14', 'lj_scale_factor',
-                    'cf_scale_factor',
-                    'LJ_type_A', 'LJ_type_B'],
-            outputs=['frc_f', 'atom_energy'])
-        self.add_prim_attr('dihedral_14_numbers', self.dihedral_14_numbers)
-        self.add_prim_attr('atom_numbers', self.atom_numbers)
-
-    def infer_shape(self, uint_crd_f_shape, ljtype_shape, charge_shape, boxlength_f_shape, a_14_shape, b_14_shape,
-                    lj_scale_factor_shape, cf_scale_factor_shape, lj_type_a_shape, lj_type_b_shape):
-        cls_name = self.name
-        n = self.atom_numbers
-        m = self.dihedral_14_numbers
-        q = lj_type_a_shape[0]
-        validator.check_int(len(uint_crd_f_shape), 2, Rel.EQ, "uint_crd_f_dim", cls_name)
-        validator.check_int(len(ljtype_shape), 1, Rel.EQ, "LJtype_dim", cls_name)
-        validator.check_int(len(charge_shape), 1, Rel.EQ, "charge_dim", cls_name)
-        validator.check_int(len(boxlength_f_shape), 1, Rel.EQ, "boxlength_f_dim", cls_name)
-        validator.check_int(len(a_14_shape), 1, Rel.EQ, "a_14_dim", cls_name)
-        validator.check_int(len(b_14_shape), 1, Rel.EQ, "b_14_dim", cls_name)
-        validator.check_int(len(lj_scale_factor_shape), 1, Rel.EQ, "lj_scale_factor_dim", cls_name)
-        validator.check_int(len(cf_scale_factor_shape), 1, Rel.EQ, "cf_scale_factor_dim", cls_name)
-        validator.check_int(len(lj_type_b_shape), 1, Rel.EQ, "LJ_type_B_dim", cls_name)
-
-        validator.check_int(uint_crd_f_shape[0], n, Rel.EQ, "uint_crd_f_shape[0]", cls_name)
-        validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape[1]", cls_name)
-        validator.check_int(ljtype_shape[0], n, Rel.EQ, "LJtype_shape", cls_name)
-        validator.check_int(charge_shape[0], n, Rel.EQ, "charge_shape", cls_name)
-        validator.check_int(boxlength_f_shape[0], 3, Rel.EQ, "boxlength_f_shape", cls_name)
-        validator.check_int(lj_type_b_shape[0], q, Rel.EQ, "LJ_type_B_shape", cls_name)
-        validator.check_int(a_14_shape[0], m, Rel.EQ, "a_14_shape", cls_name)
-        validator.check_int(b_14_shape[0], m, Rel.EQ, "b_14_shape", cls_name)
-        validator.check_int(lj_scale_factor_shape[0], m, Rel.EQ, "lj_scale_factor_shape", cls_name)
-        validator.check_int(cf_scale_factor_shape[0], m, Rel.EQ, "cf_scale_factor_shape", cls_name)
-        return uint_crd_f_shape, charge_shape
-
-    def infer_dtype(self, uint_crd_f_dtype, ljtype_dtype, charge_dtype, boxlength_f_type, a_14_type, b_14_type,
-                    lj_scale_factor_type, cf_scale_factor_type, lj_type_a_type, lj_type_b_type):
-        validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name)
-        validator.check_tensor_dtype_valid('LJtype', ljtype_dtype, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('charge', charge_dtype, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('boxlength_f', boxlength_f_type, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('a_14', a_14_type, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('b_14', b_14_type, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('lj_scale_factor', lj_scale_factor_type, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('cf_scale_factor', cf_scale_factor_type, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('LJ_type_A', lj_type_a_type, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('LJ_type_B', lj_type_b_type, [mstype.float32], self.name)
-
-        return charge_dtype, charge_dtype
-
-
-class Dihedral14LJAtomEnergy(PrimitiveWithInfer):
-    """
-    Add the potential energy caused by Lennard-Jones energy correction for each
-    necessary dihedral 1,4 terms to the total potential energy of each atom.
-
-    The calculation formula is the same as operator Dihedral14LJEnergy().
-
-    Because there is a large amount of inputs and each of them are related,
-    there is no way to construct `Examples` using random methods. For details, refer the webpage `SPONGE in MindSpore
-    <https://gitee.com/mindspore/mindscience/blob/master/MindSPONGE/docs/simple_formula.md>`_.
-
-    Args:
-        nb14_numbers (int32): the number of necessary dihedral 1,4 terms m.
-        atom_numbers (int32): the number of atoms n.
-
-    Inputs:
-        - **uint_crd_f** (Tensor) - The unsigned int coordinate value of each atom.
-          The data type is uint32 and the shape is :math:`(n, 3)`.
-        - **LJ_type** (Tensor) - The Lennard-Jones type of each atom.
-          The data type is int32 and the shape is :math:`(n,)`.
-        - **charge** (Tensor) - The charge of each atom.
-          The data type is float32 and the shape is :math:`(n,)`.
-        - **boxlength_f** (Tensor) - The length of molecular simulation box in 3 dimensions.
-          The data type is float32 and the shape is :math:`(3,)`.
-        - **a_14** (Tensor) - The first atom index of each dihedral 1,4 term.
-          The data type is int32 and the shape is :math:`(m,)`.
-        - **b_14** (Tensor) - The second atom index of each dihedral 1,4 term.
-          The data type is int32 and the shape is :math:`(m,)`.
-        - **lj_scale_factor** (Tensor) - The scale factor for the
-          Lennard-Jones part of force correction of each dihedral 1,4 term.
-          The data type is float32 and the shape is :math:`(m,)`.
-        - **cf_scale_factor** (Tensor) - The scale factor for the
-          Coulomb part of force correction for each dihedral 1,4 terms.
-          The data type is float32 and the shape is :math:`(m,)`.
-        - **LJ_type_A** (Tensor) - The A parameter in Lennard-Jones scheme of each atom pair type.
-          q is the number of atom pair. The data type is float32 and the shape is :math:`(q,)`.
-        - **LJ_type_B** (Tensor) - The B parameter in Lennard-Jones scheme of each atom pair type.
-          q is the number of atom pair. The data type is float32 and the shape is :math:`(q,)`.
-
-    Outputs:
-        - **ene** (Tensor) - The accumulated potential energy of each atom.
-          The data type is float32 and the shape is :math:`(n,)`.
-
-    Supported Platforms:
-        ``GPU``
-    """
-
-    @prim_attr_register
-    def __init__(self, nb14_numbers, atom_numbers):
-        """Initialize Dihedral14LJAtomEnergy."""
-        validator.check_value_type('nb14_numbers', nb14_numbers, int, self.name)
-        validator.check_value_type('atom_numbers', atom_numbers, int, self.name)
-        self.dihedral_14_numbers = nb14_numbers
-        self.atom_numbers = atom_numbers
-
-        self.init_prim_io_names(
-            inputs=['uint_crd_f', 'LJtype', 'charge', 'boxlength_f', 'a_14', 'b_14', 'lj_scale_factor',
-                    'LJ_type_A', 'LJ_type_B'],
-            outputs=['ene'])
-        self.add_prim_attr('dihedral_14_numbers', self.dihedral_14_numbers)
-        self.add_prim_attr('atom_numbers', self.atom_numbers)
-
-    def infer_shape(self, uint_crd_f_shape, ljtype_shape, charge_shape, boxlength_f_shape, a_14_shape, b_14_shape,
-                    lj_scale_factor_shape, lj_type_a_shape, lj_type_b_shape):
-        cls_name = self.name
-        n = self.atom_numbers
-        q = lj_type_a_shape[0]
-        validator.check_int(len(uint_crd_f_shape), 2, Rel.EQ, "uint_crd_f_dim", cls_name)
-        validator.check_int(len(ljtype_shape), 1, Rel.EQ, "LJtype_dim", cls_name)
-        validator.check_int(len(charge_shape), 1, Rel.EQ, "charge_dim", cls_name)
-        validator.check_int(len(boxlength_f_shape), 1, Rel.EQ, "boxlength_f_dim", cls_name)
-        validator.check_int(len(a_14_shape), 1, Rel.EQ, "a_14_dim", cls_name)
-        validator.check_int(len(b_14_shape), 1, Rel.EQ, "b_14_dim", cls_name)
-        validator.check_int(len(lj_scale_factor_shape), 1, Rel.EQ, "lj_scale_factor_dim", cls_name)
-        validator.check_int(len(lj_type_b_shape), 1, Rel.EQ, "LJ_type_B_dim", cls_name)
-
-        validator.check_int(uint_crd_f_shape[0], n, Rel.EQ, "uint_crd_f_shape[0]", cls_name)
-        validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape[1]", cls_name)
-        validator.check_int(ljtype_shape[0], n, Rel.EQ, "LJtype_shape", cls_name)
-        validator.check_int(charge_shape[0], n, Rel.EQ, "charge_shape", cls_name)
-        validator.check_int(boxlength_f_shape[0], 3, Rel.EQ, "boxlength_f_shape", cls_name)
-        validator.check_int(lj_type_b_shape[0], q, Rel.EQ, "LJ_type_B_shape", cls_name)
-        m = self.dihedral_14_numbers
-        validator.check_int(a_14_shape[0], m, Rel.EQ, "a_14_shape", cls_name)
-        validator.check_int(b_14_shape[0], m, Rel.EQ, "b_14_shape", cls_name)
-        validator.check_int(lj_scale_factor_shape[0], m, Rel.EQ, "lj_scale_factor_shape", cls_name)
-        return ljtype_shape
-
-    def infer_dtype(self, uint_crd_f_dtype, ljtype_dtype, charge_dtype, boxlength_f_type, a_14_type, b_14_type,
-                    lj_scale_factor_type, lj_type_a_type, lj_type_b_type):
-        validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name)
-        validator.check_tensor_dtype_valid('LJtype', ljtype_dtype, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('charge', charge_dtype, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('boxlength_f', boxlength_f_type, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('a_14', a_14_type, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('b_14', b_14_type, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('lj_scale_factor', lj_scale_factor_type, [mstype.float32],
-                                           self.name)
-        validator.check_tensor_dtype_valid('LJ_type_A', lj_type_a_type, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('LJ_type_B', lj_type_b_type, [mstype.float32], self.name)
-
-        return lj_type_a_type
-
-
-class Dihedral14CFEnergy(PrimitiveWithInfer):
-    """
-    Calculate the Coulumb part of 1,4 dihedral energy correction for
-    each necessary dihedral terms on the corresponding atoms.
-
-    Because there is a large amount of inputs and each of them are related,
-    there is no way to construct `Examples` using random methods. For details, refer the webpage `SPONGE in MindSpore
-    <https://gitee.com/mindspore/mindscience/blob/master/MindSPONGE/docs/simple_formula.md>`_.
-
-    .. math::
-
-        dr = (x_a-x_b, y_a-y_b, z_a-z_b)
-
-    .. math::
-        E = k*q_a*q_b/|dr|
-
-    Args:
-        nb14_numbers (int32): the number of necessary dihedral 1,4 terms m.
-        atom_numbers (int32): the number of atoms n.
-
-    Inputs:
-        - **uint_crd_f** (Tensor) - The unsigned int coordinate value of each atom.
-          The data type is uint32 and the shape is :math:`(n, 3)`.
-        - **LJ_type** (Tensor) - The Lennard-Jones type of each atom.
-          The data type is int32 and the shape is :math:`(n,)`.
-        - **charge** (Tensor) - The charge of each atom.
-          The data type is float32 and the shape is :math:`(n,)`.
-        - **boxlength_f** (Tensor) - The length of molecular simulation box in 3 dimensions.
-          The data type is float32 and the shape is :math:`(3,)`.
-        - **a_14** (Tensor) - The first atom index of each dihedral 1,4 term.
-          The data type is int32 and the shape is :math:`(m,)`.
-        - **b_14** (Tensor) - The second atom index of each dihedral 1,4 term.
-          The data type is int32 and the shape is :math:`(m,)`.
-        - **cf_scale_factor** (Tensor) - The scale factor for the
-          Coulomb part of force correction for each dihedral 1,4 terms.
-          The data type is float32 and the shape is :math:`(m,)`.
-
-    Outputs:
-        - **ene** (Tensor) - The accumulated potential energy of each atom.
-          The data type is float32 and the shape is :math:`(m,)`.
-
-    Supported Platforms:
-        ``GPU``
-    """
-
-    @prim_attr_register
-    def __init__(self, nb14_numbers, atom_numbers):
-        """Initialize Dihedral14CFEnergy."""
-        validator.check_value_type('nb14_numbers', nb14_numbers, int, self.name)
-        validator.check_value_type('atom_numbers', atom_numbers, int, self.name)
-        self.dihedral_14_numbers = nb14_numbers
-        self.atom_numbers = atom_numbers
-
-        self.init_prim_io_names(
-            inputs=['uint_crd_f', 'LJtype', 'charge', 'boxlength_f', 'a_14', 'b_14', 'cj_scale_factor'],
-            outputs=['ene'])
-        self.add_prim_attr('dihedral_14_numbers', self.dihedral_14_numbers)
-        self.add_prim_attr('atom_numbers', self.atom_numbers)
-
-    def infer_shape(self, uint_crd_f_shape, ljtype_shape, charge_shape, boxlength_f_shape, a_14_shape, b_14_shape,
-                    cf_scale_factor_shape):
-        cls_name = self.name
-        n = self.atom_numbers
-        validator.check_int(len(uint_crd_f_shape), 2, Rel.EQ, "uint_crd_f_dim", cls_name)
-        validator.check_int(len(ljtype_shape), 1, Rel.EQ, "LJtype_dim", cls_name)
-        validator.check_int(len(charge_shape), 1, Rel.EQ, "charge_dim", cls_name)
-        validator.check_int(len(boxlength_f_shape), 1, Rel.EQ, "boxlength_f_dim", cls_name)
-        validator.check_int(len(a_14_shape), 1, Rel.EQ, "a_14_dim", cls_name)
-        validator.check_int(len(b_14_shape), 1, Rel.EQ, "b_14_dim", cls_name)
-        validator.check_int(len(cf_scale_factor_shape), 1, Rel.EQ, "cf_scale_factor_dim", cls_name)
-
-        validator.check_int(uint_crd_f_shape[0], n, Rel.EQ, "uint_crd_f_shape[0]", cls_name)
-        validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape[1]", cls_name)
-        validator.check_int(ljtype_shape[0], n, Rel.EQ, "LJtype_shape", cls_name)
-        validator.check_int(charge_shape[0], n, Rel.EQ, "charge_shape", cls_name)
-        validator.check_int(boxlength_f_shape[0], 3, Rel.EQ, "boxlength_f_shape", cls_name)
-        m = self.dihedral_14_numbers
-        validator.check_int(a_14_shape[0], m, Rel.EQ, "a_14_shape", cls_name)
-        validator.check_int(b_14_shape[0], m, Rel.EQ, "b_14_shape", cls_name)
-        validator.check_int(cf_scale_factor_shape[0], m, Rel.EQ, "cf_scale_factor_shape", cls_name)
-        return [self.dihedral_14_numbers,]
-
-    def infer_dtype(self, uint_crd_f_dtype, ljtype_dtype, charge_dtype, boxlength_f_type, a_14_type, b_14_type,
-                    cf_scale_factor_type):
-        validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name)
-        validator.check_tensor_dtype_valid('LJtype', ljtype_dtype, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('charge', charge_dtype, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('boxlength_f', boxlength_f_type, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('a_14', a_14_type, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('b_14', b_14_type, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('lj_scale_factor', cf_scale_factor_type, [mstype.float32],
-                                           self.name)
-
-        return charge_dtype
-
-
-class Dihedral14CFAtomEnergy(PrimitiveWithInfer):
-    """
-    Add the potential energy caused by Coulumb energy correction for each
-    necessary dihedral 1,4 terms to the total potential energy of each atom.
-
-    The calculation formula is the same as operator :class:`Dihedral14CFEnergy`.
-
-    Because there is a large amount of inputs and each of them are related,
-    there is no way to construct `Examples` using random methods. For details, refer the webpage `SPONGE in MindSpore
-    <https://gitee.com/mindspore/mindscience/blob/master/MindSPONGE/docs/simple_formula.md>`_.
-
-    Args:
-        nb14_numbers (int32): the number of necessary dihedral 1,4 terms m.
-        atom_numbers (int32): the number of atoms n.
-
-    Inputs:
-        - **uint_crd_f** (Tensor) - The unsigned int coordinate value of each atom.
-          The data type is uint32 and the shape is :math:`(n, 3)`.
-        - **LJ_type** (Tensor) - The Lennard-Jones type of each atom.
-          The data type is int32 and the shape is :math:`(n,)`.
-        - **charge** (Tensor) - The charge of each atom.
-          The data type is float32 and the shape is :math:`(n,)`.
-        - **boxlength_f** (Tensor) - The length of molecular simulation box in 3 dimensions.
-          The data type is float32 and the shape is :math:`(3,)`.
-        - **a_14** (Tensor) - The first atom index of each dihedral 1,4 term.
-          The data type is int32 and the shape is :math:`(m,)`.
-        - **b_14** (Tensor) - The second atom index of each dihedral 1,4 term.
-          The data type is int32 and the shape is :math:`(m,)`.
-        - **cf_scale_factor** (Tensor) - The scale factor for the
-          Coulomb part of force correction for each dihedral 1,4 terms.
-          The data type is float32 and the shape is :math:`(m,)`.
-
-    Outputs:
-        - **ene** (Tensor) - The accumulated potential energy of each atom.
-          The data type is float32 and the shape is :math:`(n,)`
-
-
-    Supported Platforms:
-        ``GPU``
-    """
-
-    @prim_attr_register
-    def __init__(self, nb14_numbers, atom_numbers):
-        """Initialize Dihedral14CFAtomEnergy."""
-        validator.check_value_type('nb14_numbers', nb14_numbers, int, self.name)
-        validator.check_value_type('atom_numbers', atom_numbers, int, self.name)
-        self.dihedral_14_numbers = nb14_numbers
-        self.atom_numbers = atom_numbers
-
-        self.init_prim_io_names(
-            inputs=['uint_crd_f', 'LJtype', 'charge', 'boxlength_f', 'a_14', 'b_14', 'cf_scale_factor'],
-            outputs=['ene'])
-        self.add_prim_attr('dihedral_14_numbers', self.dihedral_14_numbers)
-        self.add_prim_attr('atom_numbers', self.atom_numbers)
-
-    def infer_shape(self, uint_crd_f_shape, ljtype_shape, charge_shape, boxlength_f_shape, a_14_shape, b_14_shape,
-                    cf_scale_factor_shape):
-        cls_name = self.name
-        n = self.atom_numbers
-        validator.check_int(len(uint_crd_f_shape), 2, Rel.EQ, "uint_crd_f_dim", cls_name)
-        validator.check_int(len(ljtype_shape), 1, Rel.EQ, "LJtype_dim", cls_name)
-        validator.check_int(len(charge_shape), 1, Rel.EQ, "charge_dim", cls_name)
-        validator.check_int(len(boxlength_f_shape), 1, Rel.EQ, "boxlength_f_dim", cls_name)
-        validator.check_int(len(a_14_shape), 1, Rel.EQ, "a_14_dim", cls_name)
-        validator.check_int(len(b_14_shape), 1, Rel.EQ, "b_14_dim", cls_name)
-        validator.check_int(len(cf_scale_factor_shape), 1, Rel.EQ, "cf_scale_factor_dim", cls_name)
-
-        validator.check_int(uint_crd_f_shape[0], n, Rel.EQ, "uint_crd_f_shape[0]", cls_name)
-        validator.check_int(uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape[1]", cls_name)
-        validator.check_int(ljtype_shape[0], n, Rel.EQ, "LJtype_shape", cls_name)
-        validator.check_int(charge_shape[0], n, Rel.EQ, "charge_shape", cls_name)
-        validator.check_int(boxlength_f_shape[0], 3, Rel.EQ, "boxlength_f_shape", cls_name)
-        m = self.dihedral_14_numbers
-        validator.check_int(a_14_shape[0], m, Rel.EQ, "a_14_shape", cls_name)
-        validator.check_int(b_14_shape[0], m, Rel.EQ, "b_14_shape", cls_name)
-        validator.check_int(cf_scale_factor_shape[0], m, Rel.EQ, "cf_scale_factor_shape", cls_name)
-        return ljtype_shape
-
-    def infer_dtype(self, uint_crd_f_dtype, ljtype_dtype, charge_dtype, boxlength_f_type, a_14_type, b_14_type,
-                    cf_scale_factor_type):
-        validator.check_tensor_dtype_valid('uint_crd_f', uint_crd_f_dtype, [mstype.uint32], self.name)
-        validator.check_tensor_dtype_valid('LJtype', ljtype_dtype, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('charge', charge_dtype, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('boxlength_f', boxlength_f_type, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('a_14', a_14_type, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('b_14', b_14_type, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('cf_scale_factor', cf_scale_factor_type, [mstype.float32],
-                                           self.name)
-
-        return charge_dtype
-
-
-class PMEReciprocalForce(PrimitiveWithInfer):
-    """
-    Calculate the reciprocal part of long-range Coulumb force using
-    PME(Particle Meshed Ewald) method. Assume the number of atoms is n.
-
-    The detailed calculation formula of PME(Particle Meshed Ewald) method
-    can be found in this paper: A Smooth Particle Mesh Ewald Method. DOI:
-    10.1063/1.470117.
-
-    Because there is a large amount of inputs and each of them are related,
-    there is no way to construct `Examples` using random methods. For details, refer the webpage `SPONGE in MindSpore
-    <https://gitee.com/mindspore/mindscience/blob/master/MindSPONGE/docs/simple_formula.md>`_.
-
-    Args:
-        atom_numbers(int32): the number of atoms, n.
-        beta(float32): the PME beta parameter, determined by the
-                       non-bond cutoff value and simulation precision tolerance.
-        fftx(int32): the number of points for Fourier transform in dimension X.
-        ffty(int32): the number of points for Fourier transform in dimension Y.
-        fftz(int32): the number of points for Fourier transform in dimension Z.
-        box_length_0(float32): the value of boxlength idx 0
-        box_length_1(float32): the value of boxlength idx 1
-        box_length_2(float32): the value of boxlength idx 2
-
-    Inputs:
-        - **uint_crd** (Tensor) - The unsigned int coordinates value of each atom.
-          The data type is uint32 and the shape is :math:`(n, 3)`
-        - **charge** (Tensor) - The charge carried by each atom.
-          The data type is float32 and the shape is :math:`(n,)`
-
-    Outputs:
-        - **force** (Tensor) - The force felt by each atom.
-          The data type is float32 and the shape is :math:`(n, 3)`
-
-    Supported Platforms:
-        ``GPU``
-    """
-
-    @prim_attr_register
-    def __init__(self, atom_numbers, beta, fftx, ffty, fftz, box_length_0, box_length_1, box_length_2):
-        """Initialize PMEReciprocalForce."""
-        validator.check_value_type('atom_numbers', atom_numbers, int, self.name)
-        validator.check_value_type('beta', beta, float, self.name)
-        validator.check_value_type('fftx', fftx, int, self.name)
-        validator.check_value_type('ffty', ffty, int, self.name)
-        validator.check_value_type('fftz', fftz, int, self.name)
-        validator.check_value_type('box_length_0', box_length_0, float, self.name)
-        validator.check_value_type('box_length_1', box_length_1, float, self.name)
-        validator.check_value_type('box_length_2', box_length_2, float, self.name)
-        self.atom_numbers = atom_numbers
-        self.beta = beta
-        self.fftx = fftx
-        self.ffty = ffty
-        self.fftz = fftz
-        self.box_length_0 = box_length_0
-        self.box_length_1 = box_length_1
-        self.box_length_2 = box_length_2
-
-        self.init_prim_io_names(inputs=['boxlength', 'uint_crd', 'charge'],
-                                outputs=['force'])
-        self.add_prim_attr('atom_numbers', self.atom_numbers)
-        self.add_prim_attr('beta', self.beta)
-        self.add_prim_attr('fftx', self.fftx)
-        self.add_prim_attr('ffty', self.ffty)
-        self.add_prim_attr('fftz', self.fftz)
-        self.add_prim_attr('box_length_0', self.box_length_0)
-        self.add_prim_attr('box_length_1', self.box_length_1)
-        self.add_prim_attr('box_length_2', self.box_length_2)
-
-    def infer_shape(self, uint_crd_shape, charge_shape):
-        cls_name = self.name
-        n = self.atom_numbers
-        validator.check_int(len(uint_crd_shape), 2, Rel.EQ, "uint_crd_dim", cls_name)
-        validator.check_int(len(charge_shape), 1, Rel.EQ, "charge_dim", cls_name)
-
-        validator.check_int(uint_crd_shape[0], n, Rel.EQ, "uint_crd_shape[0]", cls_name)
-        validator.check_int(uint_crd_shape[1], 3, Rel.EQ, "uint_crd_shape[1]", cls_name)
-        validator.check_int(charge_shape[0], n, Rel.EQ, "charge_shape", cls_name)
-        return uint_crd_shape
-
-    def infer_dtype(self, uint_crd_type, charge_type):
-        validator.check_tensor_dtype_valid('uint_crd', uint_crd_type, [mstype.uint32], self.name)
-        validator.check_tensor_dtype_valid('charge', charge_type, [mstype.float32], self.name)
-        return charge_type
-
-
-class PMEExcludedForce(PrimitiveWithInfer):
-    """
-    Calculate the excluded  part of long-range Coulumb force using
-    PME(Particle Meshed Ewald) method. Assume the number of atoms is
-    n, and the length of excluded list is E.
-
-    Because there is a large amount of inputs and each of them are related,
-    there is no way to construct `Examples` using random methods. For details, refer the webpage `SPONGE in MindSpore
-    <https://gitee.com/mindspore/mindscience/blob/master/MindSPONGE/docs/simple_formula.md>`_.
-
-    Args:
-        atom_numbers(int32): the number of atoms, n.
-        excluded_numbers(int32): the length of excluded list, E.
-        beta(float32): the PME beta parameter, determined by the
-          non-bond cutoff value and simulation precision tolerance.
-
-    Inputs:
-        - **uint_crd** (Tensor) - The unsigned int coordinates value of each atom.
-          The data type is uint32 and the shape is :math:`(n, 3)`
-        - **scaler** (Tensor) - The scale factor between real space
-          coordinates and its unsigned int value. The data type is float32 and the shape is :math:`(3,)`
-        - **charge** (Tensor) - The charge carried by each atom.
-          The data type is float32 and the shape is :math:`(n,)`
-        - **excluded_list_start** (Tensor) - The start excluded index
-          in excluded list for each atom. The data type is int32 and the shape is :math:`(n,)`
-        - **excluded_list** (Tensor) - The contiguous join of excluded
-          list of each atom. E is the number of excluded atoms. The data type is int32 and the shape is :math:`(E,)`
-        - **excluded_atom_numbers** (Tensor) - The number of atom excluded
-          in excluded list for each atom. The data type is int32 and the shape is :math:`(n,)`
-
-    Outputs:
-        - **force** (Tensor) - The force felt by each atom.
-          The data type is float32 and the shape is :math:`(n, 3)`
-
-    Supported Platforms:
-        ``GPU``
-    """
-
-    @prim_attr_register
-    def __init__(self, atom_numbers, excluded_numbers, beta):
-        """Initialize PMEExcludedForce."""
-        validator.check_value_type('atom_numbers', atom_numbers, int, self.name)
-        validator.check_value_type('excluded_numbers', excluded_numbers, int, self.name)
-        validator.check_value_type('beta', beta, float, self.name)
-        self.atom_numbers = atom_numbers
-        self.excluded_numbers = excluded_numbers
-        self.beta = beta
-        self.init_prim_io_names(
-            inputs=['uint_crd', 'sacler', 'charge', 'excluded_list_start', 'excluded_list', 'excluded_atom_numbers'],
-            outputs=['force'])
-        self.add_prim_attr('atom_numbers', self.atom_numbers)
-        self.add_prim_attr('excluded_numbers', self.excluded_numbers)
-        self.add_prim_attr('beta', self.beta)
-
-    def infer_shape(self, uint_crd_shape, sacler_shape, charge_shape, excluded_list_start_shape, excluded_list_shape,
-                    excluded_atom_numbers_shape):
-        cls_name = self.name
-        n = self.atom_numbers
-        validator.check_int(len(uint_crd_shape), 2, Rel.EQ, "uint_crd_dim", cls_name)
-        validator.check_int(len(sacler_shape), 1, Rel.EQ, "sacler_dim", cls_name)
-        validator.check_int(len(charge_shape), 1, Rel.EQ, "charge_dim", cls_name)
-        validator.check_int(len(excluded_list_start_shape), 1, Rel.EQ, "excluded_list_start_dim", cls_name)
-        validator.check_int(len(excluded_atom_numbers_shape), 1, Rel.EQ, "excluded_atom_numbers_dim", cls_name)
-
-        validator.check_int(uint_crd_shape[0], n, Rel.EQ, "uint_crd_shape[0]", cls_name)
-        validator.check_int(uint_crd_shape[1], 3, Rel.EQ, "uint_crd_shape[1]", cls_name)
-        validator.check_int(sacler_shape[0], 3, Rel.EQ, "sacler_shape", cls_name)
-        validator.check_int(charge_shape[0], n, Rel.EQ, "charge_shape", cls_name)
-        validator.check_int(excluded_list_start_shape[0], n, Rel.EQ, "excluded_list_start_shape", cls_name)
-        validator.check_int(excluded_atom_numbers_shape[0], n, Rel.EQ, "excluded_atom_numbers_shape", cls_name)
-        return uint_crd_shape
-
-    def infer_dtype(self, uint_crd_type, sacler_type, charge_type, excluded_list_start_type, excluded_list_type,
-                    excluded_atom_numbers_type):
-        validator.check_tensor_dtype_valid('sacler', sacler_type, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('uint_crd', uint_crd_type, [mstype.uint32], self.name)
-        validator.check_tensor_dtype_valid('charge', charge_type, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('excluded_list_start', excluded_list_start_type, [mstype.int32],
-                                           self.name)
-        validator.check_tensor_dtype_valid('excluded_list', excluded_list_type, [mstype.int32],
-                                           self.name)
-        validator.check_tensor_dtype_valid('excluded_atom_numbers', excluded_atom_numbers_type, [mstype.int32],
-                                           self.name)
-        return charge_type
-
-
-class PMEEnergy(PrimitiveWithInfer):
-    """
-    Calculate the Coulumb energy of the system using PME method.
-
-    Because there is a large amount of inputs and each of them are related,
-    there is no way to construct `Examples` using random methods. For details, refer the webpage `SPONGE in MindSpore
-    <https://gitee.com/mindspore/mindscience/blob/master/MindSPONGE/docs/simple_formula.md>`_.
-
-    .. math::
-
-        E = sum_{ij} q_iq_j/r_{ij}
-
-    Args:
-        atom_numbers(int32): the number of atoms, n.
-        excluded_numbers(int32): the length of excluded list, E.
-        beta(float32): the PME beta parameter, determined by the
-                       non-bond cutoff value and simulation precision tolerance.
-        fftx(int32): the number of points for Fourier transform in dimension X.
-        ffty(int32): the number of points for Fourier transform in dimension Y.
-        fftz(int32): the number of points for Fourier transform in dimension Z.
-        box_length_0(float32): the value of boxlength idx 0
-        box_length_1(float32): the value of boxlength idx 1
-        box_length_2(float32): the value of boxlength idx 2
-
-
-    Inputs:
-        - **uint_crd** (Tensor) - The unsigned int coordinates value of each atom.
-          The data type is uint32 and the shape is :math:`(n, 3)`
-        - **charge** (Tensor) - The charge carried by each atom.
-          The data type is float32 and the shape is :math:`(n,)`
-        - **nl_numbers** - (Tensor) - The each atom.
-          The data type is int32 and the shape is :math:`(n, 3)`
-        - **nl_serial** - (Tensor) - The neighbor list of each atom, the max number is 800.
-          The data type is int32 and the shape is :math:`(n, 800)`
-        - **scaler** (Tensor) - The scale factor between real space
-          coordinates and its unsigned int value. The data type is float32 and the shape is :math:`(3,)`
-        - **excluded_list_start** (Tensor) - The start excluded index
-          in excluded list for each atom. The data type is int32 and the shape is :math:`(n,)`
-        - **excluded_list** (Tensor) - The contiguous join of excluded
-          list of each atom. E is the number of excluded atoms. The data type is int32 and the shape is :math:`(E,)`
-        - **excluded_atom_numbers** (Tensor) - The number of atom excluded
-          in excluded list for each atom. The data type is int32 and the shape is :math:`(n,)`
-
-    Outputs:
-        - **reciprocal_ene** (Tensor) - The reciprocal term of PME energy.
-          The data type is float32 and the the shape is :math:`(1,)`.
-        - **self_ene** (Tensor) - The self term of PME energy.
-          The data type is float32 and the the shape is :math:`(1,)`.
-        - **direct_ene** (Tensor) - The direct term of PME energy.
-          The data type is float32 and the the shape is :math:`(1,)`.
-        - **correction_ene** (Tensor) - The correction term of PME energy.
-          The data type is float32 and the the shape is :math:`(1,)`.
-
-    Supported Platforms:
-        ``GPU``
-    """
-
-    @prim_attr_register
-    def __init__(self, atom_numbers, excluded_numbers, beta, fftx, ffty, fftz, box_length_0, box_length_1,
-                 box_length_2):
-        """Initialize PMEEnergy."""
-        validator.check_value_type('atom_numbers', atom_numbers, int, self.name)
-        validator.check_value_type('excluded_numbers', excluded_numbers, int, self.name)
-        validator.check_value_type('beta', beta, float, self.name)
-        validator.check_value_type('fftx', fftx, int, self.name)
-        validator.check_value_type('ffty', ffty, int, self.name)
-        validator.check_value_type('fftz', fftz, int, self.name)
-        validator.check_value_type('box_length_0', box_length_0, float, self.name)
-        validator.check_value_type('box_length_1', box_length_1, float, self.name)
-        validator.check_value_type('box_length_2', box_length_2, float, self.name)
-        self.atom_numbers = atom_numbers
-        self.excluded_numbers = excluded_numbers
-        self.beta = beta
-        self.fftx = fftx
-        self.ffty = ffty
-        self.fftz = fftz
-        self.box_length_0 = box_length_0
-        self.box_length_1 = box_length_1
-        self.box_length_2 = box_length_2
-        self.init_prim_io_names(
-            inputs=['box_length', 'uint_crd', 'charge', 'nl_numbers', 'nl_serial', 'scaler', 'excluded_list_start',
-                    'excluded_list', 'excluded_atom_numbers'],
-            outputs=['reciprocal_ene', 'self_ene', 'direct_ene', 'correction_ene'])
-        self.add_prim_attr('atom_numbers', self.atom_numbers)
-        self.add_prim_attr('excluded_numbers', self.excluded_numbers)
-        self.add_prim_attr('beta', self.beta)
-        self.add_prim_attr('fftx', self.fftx)
-        self.add_prim_attr('ffty', self.ffty)
-        self.add_prim_attr('fftz', self.fftz)
-        self.add_prim_attr('box_length_0', self.box_length_0)
-        self.add_prim_attr('box_length_1', self.box_length_1)
-        self.add_prim_attr('box_length_2', self.box_length_2)
-
-    def infer_shape(self, uint_crd, charge, nl_numbers, nl_serial, scaler, excluded_list_start,
-                    excluded_list, excluded_atom_numbers):
-        cls_name = self.name
-        n = self.atom_numbers
-        validator.check_int(len(uint_crd), 2, Rel.EQ, "uint_crd_dim", cls_name)
-        validator.check_int(len(charge), 1, Rel.EQ, "charge_dim", cls_name)
-        validator.check_int(len(nl_numbers), 1, Rel.EQ, "nl_numbers_dim", cls_name)
-        validator.check_int(len(nl_serial), 2, Rel.LE, "nl_serial_dim", cls_name)
-        validator.check_int(len(excluded_list_start), 1, Rel.EQ, "excluded_list_start_dim", cls_name)
-        validator.check_int(len(excluded_atom_numbers), 1, Rel.EQ, "excluded_atom_numbers_dim", cls_name)
-        validator.check_int(len(excluded_list), 1, Rel.GE, "excluded_list", cls_name)
-
-        validator.check_int(uint_crd[0], n, Rel.EQ, "uint_crd_shape[0]", cls_name)
-        validator.check_int(uint_crd[1], 3, Rel.EQ, "uint_crd_shape[1]", cls_name)
-        validator.check_int(charge[0], n, Rel.EQ, "charge_shape", cls_name)
-        validator.check_int(nl_numbers[0], n, Rel.EQ, "nl_numbers_shape[0]", cls_name)
-        validator.check_int(nl_serial[0], n, Rel.LE, "nl_serial_shape[0]", cls_name)
-        validator.check_int(nl_serial[1], 800, Rel.LE, "nl_serial_shape[1]", cls_name)
-        validator.check_int(excluded_list_start[0], n, Rel.EQ, "excluded_list_start_shape", cls_name)
-        validator.check_int(excluded_atom_numbers[0], n, Rel.EQ, "excluded_atom_numbers_shape", cls_name)
-        validator.check_int(excluded_list[0], 0, Rel.GE, "excluded_list_shape", cls_name)
-        return (1,), (1,), (1,), (1,)
-
-    def infer_dtype(self, uint_crd, charge, nl_numbers, nl_serial, scaler, excluded_list_start,
-                    excluded_list, excluded_atom_numbers):
-        validator.check_tensor_dtype_valid('uint_crd', uint_crd, [mstype.uint32], self.name)
-        validator.check_tensor_dtype_valid('charge', charge, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('nl_numbers', nl_numbers, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('nl_serial', nl_serial, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('scaler', scaler, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('excluded_list_start', excluded_list_start, [mstype.int32],
-                                           self.name)
-        validator.check_tensor_dtype_valid('excluded_list', excluded_list, [mstype.int32],
-                                           self.name)
-        validator.check_tensor_dtype_valid('excluded_atom_numbers', excluded_atom_numbers, [mstype.int32],
-                                           self.name)
-        return charge, charge, charge, charge
-
-
-class LJEnergy(PrimitiveWithInfer):
-    """
-    Calculate the Van der Waals interaction energy described by Lennard-Jones
-    potential for each atom. Assume the number of atoms is n, and the number
-    of Lennard-Jones types for all atoms is P, which means there will be
-    q = P*(P+1)/2 types of possible Lennard-Jones interactions for all kinds
-    of atom pairs.
-
-    Because there is a large amount of inputs and each of them are related,
-    there is no way to construct `Examples` using random methods. For details, refer the webpage `SPONGE in MindSpore
-    <https://gitee.com/mindspore/mindscience/blob/master/MindSPONGE/docs/simple_formula.md>`_.
-
-    .. math::
-
-        dr = (x_a-x_b, y_a-y_b, z_a-z_b)
-
-    .. math::
-        E = A/|dr|^{12} - B/|dr|^{6}
-
-    Args:
-        atom_numbers(int32): the number of atoms, n.
-        cutoff_square(float32): the square value of cutoff.
-
-    Inputs:
-        - **uint_crd** (Tensor) - The unsigned int coordinate value of each atom.
-          The data type is uint32 and the shape is :math:`(n, 3)`
-        - **LJtype** (Tensor) - The Lennard-Jones type of each atom.
-          The data type is int32 and the shape is :math:`(n,)`
-        - **charge** (Tensor) - The charge carried by each atom.
-          The data type is float32 and the shape is :math:`(n,)`
-        - **scaler** (Tensor) - The scale factor between real
-          space coordinate and its unsigned int value. The data type is float32 and the shape is :math:`(3,)`
-        - **nl_numbers** - (Tensor) - The each atom.
-          The data type is int32 and the shape is :math:`(n,)`
-        - **nl_serial** - (Tensor) - The neighbor list of each atom, the max number is 800.
-          The data type is int32 and the shape is :math:`(n, 800)`.
-        - **d_LJ_A** (Tensor) - The Lennard-Jones A coefficient of each kind of atom pair.
-          q is the number of atom pair. The data type is float32 and the shape is :math:`(q,)`.
-        - **d_LJ_B** (Tensor) - The Lennard-Jones B coefficient of each kind of atom pair.
-          q is the number of atom pair. The data type is float32 and the shape is :math:`(q,)`.
-
-    Outputs:
-        - **d_LJ_energy_atom** (Tensor) - The Lennard-Jones potential energy of each atom.
-          The data type is float32 and the shape is :math:`(n,)`.
-        - **d_LJ_energy_sum** (Scalar), the sum of Lennard-Jones potential energy of each atom.
-          The data type is float32.
-
-    Supported Platforms:
-        ``GPU``
-    """
-
-    @prim_attr_register
-    def __init__(self, atom_numbers, cutoff_square):
-        """Initialize LJEnergy."""
-        validator.check_value_type('atom_numbers', atom_numbers, int, self.name)
-        validator.check_value_type('cutoff_square', cutoff_square, float, self.name)
-        self.atom_numbers = atom_numbers
-        self.cutoff_square = cutoff_square
-        self.init_prim_io_names(
-            inputs=['uint_crd', 'LJtype', 'charge', 'scaler', 'nl_numbers', 'nl_serial', 'd_LJ_A', 'd_LJ_B'],
-            outputs=['d_LJ_energy_atom'])
-        self.add_prim_attr('atom_numbers', self.atom_numbers)
-        self.add_prim_attr('cutoff_square', self.cutoff_square)
-
-    def infer_shape(self, uint_crd, ljtype, charge, scaler, nl_numbers, nl_serial, d_lj_a, d_lj_b):
-        cls_name = self.name
-        n = self.atom_numbers
-        q = d_lj_a[0]
-        validator.check_int(len(uint_crd), 2, Rel.EQ, "uint_crd_dim", cls_name)
-        validator.check_int(len(ljtype), 1, Rel.EQ, "LJtype_dim", cls_name)
-        validator.check_int(len(charge), 1, Rel.EQ, "charge_dim", cls_name)
-        validator.check_int(len(scaler), 1, Rel.EQ, "scaler_dim", cls_name)
-        validator.check_int(len(nl_numbers), 1, Rel.EQ, "nl_numbers_dim", cls_name)
-        validator.check_int(len(nl_serial), 2, Rel.EQ, "nl_serial_dim", cls_name)
-        validator.check_int(len(d_lj_b), 1, Rel.EQ, "d_LJ_B_dim", cls_name)
-
-        validator.check_int(uint_crd[0], n, Rel.EQ, "uint_crd_shape[0]", cls_name)
-        validator.check_int(uint_crd[1], 3, Rel.EQ, "uint_crd_shape[1]", cls_name)
-        validator.check_int(ljtype[0], n, Rel.EQ, "LJtype_shape", cls_name)
-        validator.check_int(charge[0], n, Rel.EQ, "charge_shape", cls_name)
-        validator.check_int(scaler[0], 3, Rel.EQ, "scaler_shape", cls_name)
-        validator.check_int(nl_numbers[0], n, Rel.EQ, "nl_numbers_shape", cls_name)
-        validator.check_int(nl_serial[0], n, Rel.EQ, "nl_serial_shape[0]", cls_name)
-        validator.check_int(nl_serial[1], 800, Rel.LE, "nl_serial_shape[1]", cls_name)
-        validator.check_int(len(d_lj_a), 1, Rel.EQ, "d_LJ_A_dim", cls_name)
-        validator.check_int(d_lj_a[0], q, Rel.EQ, "d_LJ_A_shape[0]", cls_name)
-        validator.check_int(d_lj_b[0], q, Rel.EQ, "d_LJ_B_shape[0]", cls_name)
-        return charge
-
-    def infer_dtype(self, uint_crd, ljtype, charge, scaler, nl_numbers, nl_serial, d_lj_a, d_lj_b):
-        validator.check_tensor_dtype_valid('uint_crd', uint_crd, [mstype.uint32], self.name)
-        validator.check_tensor_dtype_valid('LJtype', ljtype, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('charge', charge, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('scaler', scaler, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('nl_numbers', nl_numbers, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('nl_serial', nl_serial, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('d_LJ_A', d_lj_a, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('d_LJ_B', d_lj_b, [mstype.float32], self.name)
-        return charge
-
-
-class LJForce(PrimitiveWithInfer):
-    """
-    Calculate the Van der Waals interaction force described by Lennard-Jones
-    potential energy for each atom.
-
-    Because there is a large amount of inputs and each of them are related,
-    there is no way to construct `Examples` using random methods. For details, refer the webpage `SPONGE in MindSpore
-    <https://gitee.com/mindspore/mindscience/blob/master/MindSPONGE/docs/simple_formula.md>`_.
-
-    .. math::
-
-        dr = (x_a-x_b, y_a-y_b, z_a-z_b)
-
-    .. math::
-
-        F = (-12*A/|dr|^{14} + 6*B/|dr|^{8}) * dr
-
-    Args:
-        atom_numbers(int32): the number of atoms, n.
-        cutoff_square(float32): the square value of cutoff.
-
-    Inputs:
-        - **uint_crd** (Tensor) - The unsigned int coordinates value of each atom.
-          The data type is uint32 and the shape is :math:`(n, 3)`
-        - **LJtype** (Tensor) - The Lennard-Jones type of each atom.
-          The data type is int32 and the shape is :math:`(n,)`
-        - **charge** (Tensor) - The charge carried by each atom.
-          The data type is float32 and the shape is :math:`(n,)`
-        - **scaler** (Tensor) - The scale factor between real space
-          coordinates and its unsigned int value. The data type is float32 and the shape is :math:`(3,)`
-        - **nl_numbers** - (Tensor) - The each atom.
-          The data type is int32 and the shape is :math:`(n,)`
-        - **nl_serial** - (Tensor) - The neighbor list of each atom, the max number is 800.
-          The data type is int32 and the shape is :math:`(n, 800)`.
-        - **d_LJ_A** (Tensor) - The Lennard-Jones A coefficient of each kind of atom pair.
-          q is the number of atom pair. The data type is float32 and the shape is :math:`(q,)`.
-        - **d_LJ_B** (Tensor) - The Lennard-Jones B coefficient of each kind of atom pair.
-          q is the number of atom pair. The data type is float32 and the shape is :math:`(q,)`.
-
-    outputs:
-        - **frc** (Tensor) - The force felt by each atom.
-          The data type is float32 and the shape is :math:`(n, 3)`.
-
-    Supported Platforms:
-        ``GPU``
-    """
-
-    @prim_attr_register
-    def __init__(self, atom_numbers, cutoff_square):
-        """Initialize LJForce."""
-        validator.check_value_type('atom_numbers', atom_numbers, int, self.name)
-        validator.check_value_type('cutoff_square', cutoff_square, float, self.name)
-        self.atom_numbers = atom_numbers
-        self.cutoff_square = cutoff_square
-        self.init_prim_io_names(
-            inputs=['uint_crd', 'LJtype', 'charge', 'scaler', 'nl_numbers', 'nl_serial', 'd_LJ_A', 'd_LJ_B'],
-            outputs=['frc'])
-        self.add_prim_attr('atom_numbers', self.atom_numbers)
-        self.add_prim_attr('cutoff_square', self.cutoff_square)
-
-    def infer_shape(self, uint_crd, ljtype, charge, scaler, nl_numbers, nl_serial, d_lj_a, d_lj_b):
-        cls_name = self.name
-        n = self.atom_numbers
-        q = d_lj_a[0]
-        validator.check_int(len(uint_crd), 2, Rel.EQ, "uint_crd_dim", cls_name)
-        validator.check_int(len(ljtype), 1, Rel.EQ, "LJtype_dim", cls_name)
-        validator.check_int(len(charge), 1, Rel.EQ, "charge_dim", cls_name)
-        validator.check_int(len(scaler), 1, Rel.EQ, "scaler_dim", cls_name)
-        validator.check_int(len(nl_numbers), 1, Rel.EQ, "nl_numbers_dim", cls_name)
-        validator.check_int(len(nl_serial), 2, Rel.EQ, "nl_serial_dim", cls_name)
-        validator.check_int(len(d_lj_b), 1, Rel.EQ, "d_LJ_B_dim", cls_name)
-
-        validator.check_int(uint_crd[0], n, Rel.EQ, "uint_crd_shape[0]", cls_name)
-        validator.check_int(uint_crd[1], 3, Rel.EQ, "uint_crd_shape[1]", cls_name)
-        validator.check_int(ljtype[0], n, Rel.EQ, "LJtype_shape", cls_name)
-        validator.check_int(charge[0], n, Rel.EQ, "charge_shape", cls_name)
-        validator.check_int(scaler[0], 3, Rel.EQ, "scaler_shape", cls_name)
-        validator.check_int(nl_numbers[0], n, Rel.EQ, "nl_numbers_shape", cls_name)
-        validator.check_int(nl_serial[0], n, Rel.EQ, "nl_serial_shape[0]", cls_name)
-        validator.check_int(nl_serial[1], 800, Rel.EQ, "nl_serial_shape[1]", cls_name)
-        validator.check_int(len(d_lj_a), 1, Rel.EQ, "d_LJ_A_dim", cls_name)
-        validator.check_int(d_lj_a[0], q, Rel.EQ, "d_LJ_A_shape[0]", cls_name)
-        validator.check_int(d_lj_b[0], q, Rel.EQ, "d_LJ_B_shape[0]", cls_name)
-        return uint_crd
-
-    def infer_dtype(self, uint_crd, ljtype, charge, scaler, nl_numbers, nl_serial, d_lj_a, d_lj_b):
-        validator.check_tensor_dtype_valid('uint_crd', uint_crd, [mstype.uint32], self.name)
-        validator.check_tensor_dtype_valid('LJtype', ljtype, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('charge', charge, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('scaler', scaler, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('nl_numbers', nl_numbers, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('nl_serial', nl_serial, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('d_LJ_A', d_lj_a, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('d_LJ_B', d_lj_b, [mstype.float32], self.name)
-        return charge
-
-
-class LJForceWithPMEDirectForce(PrimitiveWithInfer):
-    """
-    Calculate the Lennard-Jones force and PME direct force together.
-
-    The calculation formula of Lennard-Jones part is the same as operator
-    LJForce(), and the PME direct part is within PME method.
-
-    Because there is a large amount of inputs and each of them are related,
-    there is no way to construct `Examples` using random methods. For details, refer the webpage `SPONGE in MindSpore
-    <https://gitee.com/mindspore/mindscience/blob/master/MindSPONGE/docs/simple_formula.md>`_.
-
-    Args:
-        atom_numbers(int32): the number of atoms, n.
-        cutoff_square(float32): the square value of cutoff.
-        pme_beta(float32): PME beta parameter, same as operator PMEReciprocalForce().
-
-    Inputs:
-        - **uint_crd** (Tensor) - The unsigned int coordinate value of each atom.
-          The data type is uint32 and the shape is :math:`(n, 3)`.
-        - **LJtype** (Tensor) - The Lennard-Jones type of each atom.
-          The data type is int32 and the shape is :math:`(n,)`.
-        - **charge** (Tensor) - The charge carried by each atom.
-          The data type is float32 and the shape is :math:`(n,)`.
-        - **scaler** (Tensor) - The scale factor between real
-          space coordinate and its unsigned int value.
-          The data type is float32 and the shape is :math:`(3,)`.
-        - **nl_numbers** - (Tensor) - The each atom.
-          The data type is int32 and the shape is :math:`(n,)`.
-        - **nl_serial** - (Tensor) - The neighbor list of each atom, the max number is 800.
-          The data type is int32 and the shape is :math:`(n, 800)`.
-        - **d_LJ_A** (Tensor) - The Lennard-Jones A coefficient of each kind of atom pair.
-          q is the number of atom pair. The data type is float32 and the shape is :math:`(q,)`.
-        - **d_LJ_B** (Tensor) - The Lennard-Jones B coefficient of each kind of atom pair.
-          q is the number of atom pair. The data type is float32 and the shape is :math:`(q,)`.
-
-    Outputs:
-        - **frc** (Tensor), The force felt by each atom.
-          The data type is float32 and the shape is :math:`(n, 3)`.
-
-    Supported Platforms:
-        ``GPU``
-    """
-
-    @prim_attr_register
-    def __init__(self, atom_numbers, cutoff, pme_beta):
-        """Initialize LJForceWithPMEDirectForce."""
-        validator.check_value_type('atom_numbers', atom_numbers, int, self.name)
-        validator.check_value_type('cutoff', cutoff, float, self.name)
-        validator.check_value_type('pme_beta', pme_beta, float, self.name)
-        self.atom_numbers = atom_numbers
-        self.cutoff = cutoff
-        self.pme_beta = pme_beta
-        self.init_prim_io_names(
-            inputs=['uint_crd', 'LJtype', 'charge', 'scaler', 'nl_numbers', 'nl_serial', 'd_LJ_A', 'd_LJ_B'],
-            outputs=['frc'])
-        self.add_prim_attr('atom_numbers', self.atom_numbers)
-        self.add_prim_attr('cutoff', self.cutoff)
-        self.add_prim_attr('pme_beta', self.pme_beta)
-
-    def infer_shape(self, uint_crd, ljtype, charge, scaler, nl_numbers, nl_serial, d_lj_a, d_lj_b):
-        cls_name = self.name
-        n = self.atom_numbers
-        q = d_lj_a[0]
-        validator.check_int(len(uint_crd), 2, Rel.EQ, "uint_crd_dim", cls_name)
-        validator.check_int(len(ljtype), 1, Rel.EQ, "LJtype_dim", cls_name)
-        validator.check_int(len(charge), 1, Rel.EQ, "charge_dim", cls_name)
-        validator.check_int(len(scaler), 1, Rel.EQ, "scaler_dim", cls_name)
-        validator.check_int(len(nl_numbers), 1, Rel.EQ, "nl_numbers_dim", cls_name)
-        validator.check_int(len(nl_serial), 2, Rel.EQ, "nl_serial_dim", cls_name)
-        validator.check_int(len(d_lj_b), 1, Rel.EQ, "d_LJ_B_dim", cls_name)
-
-        validator.check_int(uint_crd[0], n, Rel.EQ, "uint_crd_shape[0]", cls_name)
-        validator.check_int(uint_crd[1], 3, Rel.EQ, "uint_crd_shape[1]", cls_name)
-        validator.check_int(ljtype[0], n, Rel.EQ, "LJtype_shape", cls_name)
-        validator.check_int(charge[0], n, Rel.EQ, "charge_shape", cls_name)
-        validator.check_int(scaler[0], 3, Rel.EQ, "scaler_shape", cls_name)
-        validator.check_int(nl_numbers[0], n, Rel.EQ, "nl_numbers_shape", cls_name)
-        validator.check_int(nl_serial[0], n, Rel.EQ, "nl_serial_shape[0]", cls_name)
-        validator.check_int(nl_serial[1], 800, Rel.EQ, "nl_serial_shape[1]", cls_name)
-        validator.check_int(len(d_lj_a), 1, Rel.EQ, "d_LJ_A_dim", cls_name)
-        validator.check_int(d_lj_a[0], q, Rel.EQ, "d_LJ_A_shape[0]", cls_name)
-        validator.check_int(d_lj_b[0], q, Rel.EQ, "d_LJ_B_shape[0]", cls_name)
-        return uint_crd
-
-    def infer_dtype(self, uint_crd, ljtype, charge, scaler, nl_numbers, nl_serial, d_lj_a, d_lj_b):
-        validator.check_tensor_dtype_valid('uint_crd', uint_crd, [mstype.uint32], self.name)
-        validator.check_tensor_dtype_valid('LJtype', ljtype, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('charge', charge, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('scaler', scaler, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('nl_numbers', nl_numbers, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('nl_serial', nl_serial, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('d_LJ_A', d_lj_a, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('d_LJ_B', d_lj_b, [mstype.float32], self.name)
-        return charge
-
-
-class MDTemperature(PrimitiveWithInfer):
-    """
-    Compute the MD temperature.
-
-    Because there is a large amount of inputs and each of them are related,
-    there is no way to construct `Examples` using random methods. For details, refer the webpage `SPONGE in MindSpore
-    <https://gitee.com/mindspore/mindscience/blob/master/MindSPONGE/docs/simple_formula.md>`_.
-
-    Args:
-        residue_numbers (int32): the number of residues m.
-        atom_numbers (int32): the number of atoms n.
-
-    Inputs:
-        - **start** (Tensor) - The start atom index of each residue.
-          The data type is int32 and the shape is :math:`(m,)`.
-        - **end** (Tensor) - The end atom index of each residue.
-          The data type is int32 and the shape is :math:`(m,)`.
-        - **atom_vel_f** (Tensor) - The velocity of each atom.
-          The data type is float32 and the shape is :math:`(n, 3)`.
-        - **atom_mass** (Tensor) - The mass of each atom.
-          The data type is float32 and the shape is :math:`(n,)`.
-
-    Outputs:
-        - **ek** (Tensor) - The temperature of each atom.
-          The data type is float32 and the shape is :math:`(n,)`.
-
-    Supported Platforms:
-        ``GPU``
-    """
-
-    @prim_attr_register
-    def __init__(self, residue_numbers, atom_numbers):
-        """Initialize MDTemperature."""
-        validator.check_value_type('residue_numbers', residue_numbers, int, self.name)
-        validator.check_value_type('atom_numbers', atom_numbers, int, self.name)
-        self.residue_numbers = residue_numbers
-        self.atom_numbers = atom_numbers
-        self.add_prim_attr('residue_numbers', self.residue_numbers)
-        self.add_prim_attr('atom_numbers', self.atom_numbers)
-        self.init_prim_io_names(
-            inputs=['start', 'end', 'atom_vel_f', 'atom_mass'],
-            outputs=['ek'])
-
-    def infer_shape(self, start_shape, end_shape, atom_vel_f_shape, atom_mass_shape):
-        cls_name = self.name
-        n = self.residue_numbers
-        m = self.atom_numbers
-        validator.check_int(len(start_shape), 1, Rel.EQ, "start", cls_name)
-        validator.check_int(start_shape[0], n, Rel.EQ, "end", cls_name)
-        validator.check_int(len(end_shape), 1, Rel.EQ, "start", cls_name)
-        validator.check_int(end_shape[0], n, Rel.EQ, "end", cls_name)
-        validator.check_int(atom_vel_f_shape[0], m, Rel.EQ, "atom_vel_f", cls_name)
-        validator.check_int(atom_vel_f_shape[1], 3, Rel.EQ, "atom_vel_f", cls_name)
-        validator.check_int(len(atom_mass_shape), 1, Rel.EQ, "atom_mass", cls_name)
-        validator.check_int(atom_mass_shape[0], m, Rel.EQ, "atom_mass", cls_name)
-        return [n,]
-
-    def infer_dtype(self, start_dtype, end_dtype, atom_vel_f_dtype, atom_mass_dtype):
-        validator.check_tensor_dtype_valid('start', start_dtype, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('end', end_dtype, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('atom_vel_f', atom_vel_f_dtype, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('atom_mass', atom_mass_dtype, [mstype.float32], self.name)
-        return atom_mass_dtype
-
-
-class MDIterationLeapFrogWithRF(PrimitiveWithInfer):
-    """
-    One step of classical leap frog algorithm to solve the finite difference
-    Hamiltonian equations of motion for certain system, using Langevin dynamics
-    with Liu's thermostat scheme. Assume the number of atoms is n and the target
-    control temperature is T.
-
-    Detailed iteration formula can be found in this paper: A unified thermostat
-    scheme for efficient configurational sampling for classical/quantum canonical
-    ensembles via molecular dynamics. DOI: 10.1063/1.4991621.
-
-    Because there is a large amount of inputs and each of them are related,
-    there is no way to construct `Examples` using random methods. For details, refer the webpage `SPONGE in MindSpore
-    <https://gitee.com/mindspore/mindscience/blob/master/MindSPONGE/docs/simple_formula.md>`_.
-
-    Inputs:
-        - **float4_numbers** (Scalar) - total length to store random numbers.
-          The data type is int32.
-        - **atom_numbers** (Scalar) - The number of atoms n.
-          The data type is int32.
-        - **dt** (Scalar) - time step for finite difference. The data type is float32.
-        - **half_dt** (Scalar) - half of time step for finite difference.
-          The data type is float32.
-        - **exp_gamma** (Scalar) - parameter in Liu's dynamic, equals
-          exp(-gamma_ln * dt), where gamma_ln is the firction factor in Langvin
-          dynamics. The data type is float32.
-        - **max_velocity** (Scalar) - The upper limit of velocity, when the
-          velocity overflows, scale it to the upper limit. The data type is float32.
-        - **is_max_velocity** (Scalar) - whether the max velocity control is
-          open or not. The data type is int32.
-        - **mass_inverse** (Tensor) - The inverse value of
-          mass of each atom. The data type is float32 and the shape is :math:`(n,)`.
-        - **sqrt_mass** (Tensor) - The inverse square root value
-          of effect mass in Liu's dynamics of each atom.
-          The data type is float32 and the shape is :math:`(n,)`.
-        - **vel** (Tensor) - The velocity of each atom.
-          The data type is float32 and the shape is :math:`(n, 3)`.
-        - **crd** (Tensor) - The coordinate of each atom.
-          The data type is float32 and the shape is :math:`(n, 3)`.
-        - **frc** (Tensor) - The force felt by each atom.
-          The data type is float32 and the shape is :math:`(n, 3)`.
-        - **acc** (Tensor) - The acceleration of each atom.
-          The data type is float32 and the shape is :math:`(n, 3)`.
-        - **random force** (Tensor) - The random forces.
-          The data type is float32 and the shape is :math:`(n, 3)`.
-
-    Outputs:
-        - **res** (Scalar) - The data type is float32.
-
-    Supported Platforms:
-        ``GPU``
-    Examples:
-    """
-
-    @prim_attr_register
-    def __init__(self, float4_numbers, atom_numbers, half_dt, dt, exp_gamma, is_max_velocity, max_velocity):
-        """Initialize MDIterationLeapFrogWithRF."""
-        validator.check_value_type('float4_numbers', float4_numbers, int, self.name)
-        validator.check_value_type('atom_numbers', atom_numbers, int, self.name)
-        validator.check_value_type('half_dt', half_dt, float, self.name)
-        validator.check_value_type('dt', dt, float, self.name)
-        validator.check_value_type('exp_gamma', exp_gamma, float, self.name)
-        validator.check_value_type('is_max_velocity', is_max_velocity, int, self.name)
-        validator.check_value_type('max_velocity', max_velocity, float, self.name)
-        self.float4_numbers = float4_numbers
-        self.atom_numbers = atom_numbers
-        self.half_dt = half_dt
-        self.dt = dt
-        self.exp_gamma = exp_gamma
-        self.is_max_velocity = is_max_velocity
-        self.max_velocity = max_velocity
-
-        self.init_prim_io_names(
-            inputs=['mass_inverse', 'sqrt_mass', 'vel_in', 'crd_in', 'frc_in', 'acc_in', 'random_force'],
-            outputs=['res'])
-        self.add_prim_attr('float4_numbers', self.float4_numbers)
-        self.add_prim_attr('atom_numbers', self.atom_numbers)
-        self.add_prim_attr('half_dt', self.half_dt)
-        self.add_prim_attr('dt', self.dt)
-        self.add_prim_attr('exp_gamma', self.exp_gamma)
-        self.add_prim_attr('is_max_velocity', self.is_max_velocity)
-        self.add_prim_attr('max_velocity', self.max_velocity)
-
-    def infer_shape(self, mass_inverse_shape, sqrt_mass_shape, vel_in_shape, crd_in_shape, frc_in_shape, acc_in_shape,
-                    random_force_shape):
-        n = self.atom_numbers
-        validator.check_int(len(mass_inverse_shape), 1, Rel.EQ, "mass_inverse", self.name)
-        validator.check_int(len(sqrt_mass_shape), 1, Rel.EQ, "mass_inverse", self.name)
-        validator.check_int(mass_inverse_shape[0], n, Rel.EQ, "mass_inverse", self.name)
-        validator.check_int(sqrt_mass_shape[0], n, Rel.EQ, "mass_inverse", self.name)
-        validator.check_int(vel_in_shape[0], n, Rel.EQ, "vel_in", self.name)
-        validator.check_int(vel_in_shape[1], 3, Rel.EQ, "vel_in", self.name)
-        validator.check_int(crd_in_shape[0], n, Rel.EQ, "crd_in", self.name)
-        validator.check_int(crd_in_shape[1], 3, Rel.EQ, "crd_in", self.name)
-        validator.check_int(frc_in_shape[0], n, Rel.EQ, "frc_in", self.name)
-        validator.check_int(frc_in_shape[1], 3, Rel.EQ, "frc_in", self.name)
-        validator.check_int(acc_in_shape[0], n, Rel.EQ, "acc_in", self.name)
-        validator.check_int(acc_in_shape[1], 3, Rel.EQ, "acc_in", self.name)
-        validator.check_int(random_force_shape[0], n, Rel.EQ, "random_force", self.name)
-        validator.check_int(random_force_shape[1], 3, Rel.EQ, "random_force", self.name)
-
-        return [1,]
-
-    def infer_dtype(self, mass_inverse_dtype, sqrt_mass_dtype, vel_in_dtype, crd_in_dtype, frc_in_dtype, acc_in_dtype,
-                    rf_dtype):
-        validator.check_tensor_dtype_valid('mass_inverse', mass_inverse_dtype, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('sqrt_mass', sqrt_mass_dtype, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('vel_in', vel_in_dtype, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('crd_in', crd_in_dtype, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('frc_in', frc_in_dtype, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('acc_in', acc_in_dtype, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('rf', rf_dtype, [mstype.float32], self.name)
-        return mstype.float32
-
-
-class MDIterationLeapFrogLiujian(PrimitiveWithInfer):
-    """
-    One step of classical leap frog algorithm to solve the finite difference
-    Hamiltonian equations of motion for certain system, using Langevin dynamics
-    with Liu's thermostat scheme. Assume the number of atoms is n and the target
-    control temperature is T.
-
-    Detailed iteration formula can be found in this paper: A unified thermostat
-    scheme for efficient configurational sampling for classical/quantum canonical
-    ensembles via molecular dynamics. DOI: 10.1063/1.4991621.
-
-    Because there is a large amount of inputs and each of them are related,
-    there is no way to construct `Examples` using random methods. For details, refer the webpage `SPONGE in MindSpore
-    <https://gitee.com/mindspore/mindscience/blob/master/MindSPONGE/docs/simple_formula.md>`_.
-
-    Args:
-        atom_numbers(int32): the number of atoms n.
-        dt(float32): time step for finite difference.
-        half_dt(float32): half of time step for finite difference.
-        exp_gamma(float32): parameter in Liu's dynamic.
-
-    Inputs:
-        - **inverse_mass** (Tensor) - The inverse value of
-          mass of each atom. The data type is float32 and the shape is :math:`(n)`.
-        - **sqrt_mass_inverse** (Tensor) - The inverse square root value
-          of effect mass in Liu's dynamics of each atom.
-          The data type is float32 and the shape is :math:`(n,)`.
-        - **vel** (Tensor) - The velocity of each atom.
-          The data type is float32 and the shape is :math:`(n, 3)`.
-        - **crd** (Tensor) - The coordinate of each atom.
-          The data type is float32 and the shape is :math:`(n, 3)`.
-        - **frc** (Tensor) - The force felt by each atom.
-          The data type is float32 and the shape is :math:`(n, 3)`.
-        - **acc** (Tensor) - The acceleration of each atom.
-          The data type is float32 and the shape is :math:`(n, 3)`.
-        - **rand_state** (Tensor) - Random state to generate
-          random force. The data type is float32 and the shape is :math:`(math.ceil(n * 3.0 / 4.0) * 16, )`.
-        - **rand_frc** (Tensor) - The random forces.
-          The data type is float32 and the shape is :math:`(n, 3)`.
-
-    Outputs:
-        - **output** (Tensor) - The output coordinates.
-          The data type is float32, and the shape is :math:`(n, 3)`.
-
-    Supported Platforms:
-        ``GPU``
-    """
-
-    @prim_attr_register
-    def __init__(self, atom_numbers, half_dt, dt, exp_gamma):
-        """Initialize MDIterationLeapFrogLiujian."""
-        self.atom_numbers = atom_numbers
-        self.half_dt = half_dt
-        self.dt = dt
-        self.exp_gamma = exp_gamma
-
-        self.add_prim_attr('atom_numbers', self.atom_numbers)
-        self.add_prim_attr('half_dt', self.half_dt)
-        self.add_prim_attr('dt', self.dt)
-        self.add_prim_attr('exp_gamma', self.exp_gamma)
-        self.init_prim_io_names(
-            inputs=['inverse_mass', 'sqrt_mass_inverse', 'vel', 'crd', 'frc', 'acc', 'rand_state', 'rand_frc'],
-            outputs=['output'])
-        self.add_prim_attr('side_effect_mem', True)
-
-    def infer_shape(self, inverse_mass, sqrt_mass_inverse, vel, crd, frc, acc, rand_state, rand_frc):
-        n = self.atom_numbers
-        validator.check_int(len(inverse_mass), 1, Rel.EQ, "inverse_mass", self.name)
-        validator.check_int(len(sqrt_mass_inverse), 1, Rel.EQ, "sqrt_mass_inverse", self.name)
-        validator.check_int(inverse_mass[0], n, Rel.EQ, "inverse_mass", self.name)
-        validator.check_int(sqrt_mass_inverse[0], n, Rel.EQ, "sqrt_mass_inverse", self.name)
-        validator.check_int(len(rand_state), 1, Rel.EQ, "rand_state_dim", self.name)
-        validator.check_int(len(rand_frc), 2, Rel.EQ, "rand_frc_dim", self.name)
-        validator.check_int(len(vel), 2, Rel.EQ, "vel_dim", self.name)
-        validator.check_int(len(crd), 2, Rel.EQ, "crd_dim", self.name)
-        validator.check_int(len(frc), 2, Rel.EQ, "frc_dim", self.name)
-        validator.check_int(len(acc), 2, Rel.EQ, "acc_dim", self.name)
-        validator.check_int(vel[0], n, Rel.EQ, "vel_shape[0]", self.name)
-        validator.check_int(vel[1], 3, Rel.EQ, "vel_shape[1]", self.name)
-        validator.check_int(crd[0], n, Rel.EQ, "crd_shape[0]", self.name)
-        validator.check_int(crd[1], 3, Rel.EQ, "crd_shape[1]", self.name)
-        validator.check_int(frc[0], n, Rel.EQ, "frc_shape[0]", self.name)
-        validator.check_int(frc[1], 3, Rel.EQ, "frc_shape[1]", self.name)
-        validator.check_int(acc[0], n, Rel.EQ, "acc_shape[0]", self.name)
-        validator.check_int(acc[1], 3, Rel.EQ, "acc_shape[1]", self.name)
-        validator.check_int(rand_frc[0], n, Rel.EQ, "rand_frc_shape[0]", self.name)
-        validator.check_int(rand_frc[1], 3, Rel.EQ, "rand_frc_shape[1]", self.name)
-        validator.check_int(rand_state[0], math.ceil(self.atom_numbers * 3 / 4.0) * 16, Rel.EQ, "rand_state", self.name)
-        return [self.atom_numbers, 3]
-
-    def infer_dtype(self, inverse_mass, sqrt_mass_inverse, vel, crd, frc, acc, rand_state, rand_frc):
-        validator.check_tensor_dtype_valid('inverse_mass', inverse_mass, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('sqrt_mass_inverse', sqrt_mass_inverse, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('vel', vel, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('crd', crd, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('frc', frc, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('acc', acc, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('rand_frc', rand_frc, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('rand_state', rand_state, [mstype.float32], self.name)
-        return mstype.float32
-
-
-class CrdToUintCrd(PrimitiveWithInfer):
-    """
-    Convert FP32 coordinate to Uint32 coordinate.
-
-    Because there is a large amount of inputs and each of them are related,
-    there is no way to construct `Examples` using random methods. For details, refer the webpage `SPONGE in MindSpore
-    <https://gitee.com/mindspore/mindscience/blob/master/MindSPONGE/docs/simple_formula.md>`_.
-
-    Args:
-        atom_numbers(int32): the number of atoms n.
-
-    Inputs:
-        - **crd_to_uint_crd_cof** (Tensor) - The scale factor
-          between the unsigned int value and the real space coordinates.
-          The data type is float32 and the shape is :math:`(3,)`.
-        - **crd** (Tensor) - The coordinate of each atom.
-          The data type is float32 and the shape is :math:`(n, 3)`.
-
-    Outputs:
-        - **output** (Scalar) - The data type is uint32.
-
-    Supported Platforms:
-        ``GPU``
-    """
-
-    @prim_attr_register
-    def __init__(self, atom_numbers):
-        """Initialize CrdToUintCrd."""
-        self.atom_numbers = atom_numbers
-        self.add_prim_attr('atom_numbers', self.atom_numbers)
-        self.init_prim_io_names(
-            inputs=['crd_to_uint_crd_cof', 'crd'],
-            outputs=['output'])
-
-    def infer_shape(self, crd_to_uint_crd_cof, crd):
-        validator.check_int(crd_to_uint_crd_cof[0], 3, Rel.EQ, "crd_to_uint_crd_cof_shape", self.name)
-        validator.check_int(crd[0], self.atom_numbers, Rel.EQ, "crd[0]", self.name)
-        validator.check_int(crd[1], 3, Rel.EQ, "crd[1]", self.name)
-        return crd
-
-    def infer_dtype(self, crd_to_uint_crd_cof, crd):
-        validator.check_tensor_dtype_valid('crd_to_uint_crd_cof', crd_to_uint_crd_cof, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('crd', crd, [mstype.float32], self.name)
-        return mstype.uint32
-
-
-class MDIterationSetupRandState(PrimitiveWithInfer):
-    """
-    Compute the random state of the iteration.
-
-    Because there is a large amount of inputs and each of them are related,
-    there is no way to construct `Examples` using random methods. For details, refer the webpage `SPONGE in MindSpore
-    <https://gitee.com/mindspore/mindscience/blob/master/MindSPONGE/docs/simple_formula.md>`_.
-
-    Args:
-        atom_numbers(int32): the number of atoms n.
-        seed(int32): random seed.
-
-    Outputs:
-        - **output** (Tensor) random state.
-          The data type is float32 and the shape is :math:`(ceil(n * 3 / 4),)`.
-
-    Supported Platforms:
-        ``GPU``
-    """
-
-    @prim_attr_register
-    def __init__(self, atom_numbers, seed):
-        """Initialize MDIterationSetupRandState."""
-        self.atom_numbers = atom_numbers
-        self.seed = seed
-        self.add_prim_attr('atom_numbers', self.atom_numbers)
-        self.add_prim_attr('seed', self.seed)
-        self.init_prim_io_names(
-            inputs=[],
-            outputs=['output'])
-
-    def infer_shape(self):
-        float4_numbers = math.ceil(self.atom_numbers * 3 / 4.0)
-        curandsize = 64 / 4
-        return [float4_numbers * int(curandsize),]
-
-    def infer_dtype(self):
-        return mstype.float32
-
-
-class TransferCrd(PrimitiveWithInfer):
-    """
-    Transfer the coordinates to angular and radial.
-
-    Because there is a large amount of inputs and each of them are related,
-    there is no way to construct `Examples` using random methods. For details, refer the webpage `SPONGE in MindSpore
-    <https://gitee.com/mindspore/mindscience/blob/master/MindSPONGE/docs/simple_formula.md>`_.
-
-    Args:
-        start_serial(int32): the index start position.
-        end_serial(int32): the index end position.
-        number(int32): the length of angular and radial.
-
-    Inputs:
-        - **crd** (Tensor) - The coordinate of each atom.
-          n is the number of atoms. The data type is float32 and the shape is :math:`(n, 3)`.
-        - **old_crd** (Tensor) - The last coordinate of each atom.
-          n is the number of atoms. The data type is float32 and the shape is :math:`(n, 3)`.
-        - **box** (Tensor) - The length of 3 dimensions of the simulation box.
-          The data type is float32 and the shape is :math:`(3,)`.
-
-    Outputs:
-        - **radial** (Tensor) - The array of radial transferred from coordinates.
-          The data type is float32 and the shape is :math:`(number,)`.
-        - **angular** (Tensor) - The array of angular transferred from coordinates.
-          The data type is float32 and the shape is :math:`(number,)`.
-        - **nowarp_crd** (Tensor) - The modified coordinate of each atom for
-          computing radial and angular. The data type is float32 and the shape is :math:`(n, 3)`.
-        - **box_map_times** (Tensor) - The box map times for radial and  angular.
-          The data type is int32 and the shape is :math:`(n, 3)`.
-
-    Supported Platforms:
-        ``GPU``
-    """
-
-    @prim_attr_register
-    def __init__(self, start_serial, end_serial, number, atom_numbers):
-        """Initialize TransferCrd."""
-        validator.check_value_type('start_serial', start_serial, int, self.name)
-        validator.check_value_type('end_serial', end_serial, int, self.name)
-        validator.check_value_type('number', number, int, self.name)
-        validator.check_value_type('atom_numbers', atom_numbers, int, self.name)
-        self.start_serial = start_serial
-        self.end_serial = end_serial
-        self.number = number
-        self.atom_numbers = atom_numbers
-        self.add_prim_attr('start_serial', self.start_serial)
-        self.add_prim_attr('end_serial', self.end_serial)
-        self.add_prim_attr('number', self.number)
-        self.add_prim_attr('atom_numbers', self.atom_numbers)
-        self.init_prim_io_names(
-            inputs=['crd', 'old_crd', 'box'],
-            outputs=['radial', 'angular', 'nowarp_crd', 'box_map_times'])
-
-    def infer_shape(self, crd_shape, old_crd_shape, box_shape):
-        n = self.atom_numbers
-        validator.check_int(len(crd_shape), 2, Rel.EQ, "crd_dim", self.name)
-        validator.check_int(crd_shape[0], n, Rel.EQ, "crd_shape[0]", self.name)
-        validator.check_int(crd_shape[1], 3, Rel.EQ, "crd_shape[1]", self.name)
-        validator.check_int(len(old_crd_shape), 2, Rel.EQ, "old_crd_dim", self.name)
-        validator.check_int(old_crd_shape[0], n, Rel.EQ, "old_crd_shape[0]", self.name)
-        validator.check_int(old_crd_shape[1], 3, Rel.EQ, "old_crd_shape[1]", self.name)
-        validator.check_int(len(box_shape), 1, Rel.EQ, "box_dim", self.name)
-        validator.check_int(box_shape[0], 3, Rel.EQ, "box_shape[0]", self.name)
-        return [self.number,], [self.number,], [n, 3], [n, 3]
-
-    def infer_dtype(self, crd_dtype, old_crd_dtype, box_dtype):
-        validator.check_tensor_dtype_valid('crd', crd_dtype, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('old_crd', old_crd_dtype, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('box', box_dtype, [mstype.float32], self.name)
-        return mstype.float32, mstype.float32, mstype.float32, mstype.int32
-
-
-class PMEBatchedFFT2D(PrimitiveWithInfer):
-    """
-    Forward FFT with N-Dimensional Input. currently this is only batched ifft2.
-
-    .. warning::
-        This is an experimental prototype that is subject to change and/or deletion.
-
-    Inputs:
-        - **input_tensor** (Tensor) - Three dimentsional tensor, supported
-          data type is complex64.
-
-    Outputs:
-        - **output_tensor** (Tensor) - The tensor after undergoing fast Fourier
-          transform, the data type is complex64.
-
-    Supported Platforms:
-        ``GPU``
-    """
-
-    @prim_attr_register
-    def __init__(self, direction):
-        self.init_prim_io_names(
-            inputs=['input_tensor'],
-            outputs=['output_tensor'])
-
-        validator.check_value_type('direction', direction, int, self.name)
-        self.direction = direction
-        self.add_prim_attr('direction', self.direction)
-
-    def infer_shape(self, input_shape):
-        self.add_prim_attr("fftx", input_shape[0])
-        self.add_prim_attr("ffty", input_shape[1])
-        self.add_prim_attr("fftz", input_shape[2])
-        return [input_shape[0], input_shape[1], input_shape[2]]
-
-    def infer_dtype(self, input_dtype):
-        validator.check_tensor_dtype_valid('input_tensor', input_dtype, [mstype.complex64], self.name)
-        return mstype.complex64
-
-
-class PMEFFT1D(PrimitiveWithInfer):
-    """
-    Forward FFT with One-Dimensional Input.
-
-    .. warning::
-        This is an experimental prototype that is subject to change and/or deletion.
-
-    Inputs:
-        - **input_tensor** (Tensor) - Three dimentsional tensor, supported
-          data type is complex64.
-
-    Outputs:
-        - **output_tensor** (Tensor) - The tensor after undergoing fast Fourier
-          transform, the data type is complex64.
-
-    Supported Platforms:
-        ``GPU``
-    """
-
-    @prim_attr_register
-    def __init__(self):
-        self.init_prim_io_names(
-            inputs=['input_tensor'],
-            outputs=['output_tensor'])
-
-    def infer_shape(self, input_shape):
-        self.add_prim_attr('fftx', input_shape[0])
-        return [input_shape[0]]
-
-    def infer_dtype(self, input_dtype):
-        validator.check_tensor_dtype_valid('input_tensor', input_dtype, [mstype.complex64], self.name)
-        return mstype.complex64
-
-
-class PMEIFFT1D(PrimitiveWithInfer):
-    """
-    Inverse FFT with One-Dimensional Input.
-
-    .. warning::
-        This is an experimental prototype that is subject to change and/or deletion.
-
-    Inputs:
-        - **input_tensor** (Tensor) - Three dimentsional input tensor, supported data
-          type is complex64.
-
-    Outputs:
-        - **output_tensor** (Tensor) - Returns the tensor after undergoing
-          inverse Fourier transform, the data type is complex64.
-
-    Supported Platforms:
-        ``GPU``
-    """
-
-    @prim_attr_register
-    def __init__(self):
-        self.init_prim_io_names(
-            inputs=['input_tensor'],
-            outputs=['output_tensor'])
-
-    def infer_shape(self, input_shape):
-        self.add_prim_attr('fftx', input_shape[0])
-        return [input_shape[0]]
-
-    def infer_dtype(self, input_dtype):
-        validator.check_tensor_dtype_valid('input_tensor', input_dtype, [mstype.complex64], self.name)
-        return mstype.complex64
-
-
-class PMEFFT2D(PrimitiveWithInfer):
-    """
-    Forward FFT with Two-Dimensional Input.
-
-    .. warning::
-        This is an experimental prototype that is subject to change and/or deletion.
-
-    Inputs:
-        - **input_tensor** (Tensor) - Three dimentsional tensor, supported
-          data type is complex64.
-
-    Outputs:
-        - **output_tensor** (Tensor) - The tensor after undergoing fast Fourier
-          transform, the data type is complex64.
-
-    Supported Platforms:
-        ``GPU``
-    """
-
-    @prim_attr_register
-    def __init__(self):
-        self.init_prim_io_names(
-            inputs=['input_tensor'],
-            outputs=['output_tensor'])
-
-    def infer_shape(self, input_shape):
-        self.add_prim_attr('fftx', input_shape[0])
-        self.add_prim_attr('ffty', input_shape[1])
-        return [input_shape[0], input_shape[1]]
-
-    def infer_dtype(self, input_dtype):
-        validator.check_tensor_dtype_valid('input_tensor', input_dtype, [mstype.complex64], self.name)
-        return mstype.complex64
-
-
-class PMEIFFT2D(PrimitiveWithInfer):
-    """
-    Inverse FFT with Two-Dimensional Input.
-
-    .. warning::
-        This is an experimental prototype that is subject to change and/or deletion.
-
-    Inputs:
-        - **input_tensor** (Tensor) - Three dimentsional input tensor, supported data
-          type is complex64.
-
-    Outputs:
-        - **output_tensor** (Tensor) - Return the tensor after undergoing
-          inverse Fourier transform, the data type is complex64.
-
-    Supported Platforms:
-        ``GPU``
-    """
-
-    @prim_attr_register
-    def __init__(self):
-        self.init_prim_io_names(
-            inputs=['input_tensor'],
-            outputs=['output_tensor'])
-
-    def infer_shape(self, input_shape):
-        self.add_prim_attr('fftx', input_shape[0])
-        self.add_prim_attr('ffty', input_shape[1])
-        return [input_shape[0], input_shape[1]]
-
-    def infer_dtype(self, input_dtype):
-        validator.check_tensor_dtype_valid('input_tensor', input_dtype, [mstype.complex64], self.name)
-        return mstype.complex64
-
-
-class PMERFFT2D(PrimitiveWithInfer):
-    """
-    Forward FFT with Two-Dimensional Input for real -> complex.
-
-    .. warning::
-        This is an experimental prototype that is subject to change and/or deletion.
-
-    Inputs:
-        - **input_tensor** (Tensor) - Three dimentsional tensor, supported
-          data type is float32.
-
-    Outputs:
-        - **output_tensor** (Tensor) - The tensor after undergoing fast Fourier
-          transform, the data type is complex64.
-
-    Supported Platforms:
-        ``GPU``
-    """
-
-    @prim_attr_register
-    def __init__(self):
-        self.init_prim_io_names(
-            inputs=['input_tensor'],
-            outputs=['output_tensor'])
-
-    def infer_shape(self, input_shape):
-        self.add_prim_attr('fftx', input_shape[0])
-        self.add_prim_attr('ffty', input_shape[1])
-        return [input_shape[0], int(input_shape[1]/2)+1]
-
-    def infer_dtype(self, input_dtype):
-        validator.check_tensor_dtype_valid('input_tensor', input_dtype, [mstype.float32], self.name)
-        return mstype.complex64
-
-
-class PMEIRFFT2D(PrimitiveWithInfer):
-    """
-    Inverse RFFT with Two-Dimensional Input.
-
-    .. warning::
-        This is an experimental prototype that is subject to change and/or deletion.
-
-    Inputs:
-        - **input_tensor** (Tensor) - Three dimentsional input tensor, supported data
-          type is complex64.
-
-    Outputs:
-        - **output_tensor** (Tensor) - Return the tensor after undergoing
-          inverse Fourier transform, the data type is float32.
-
-    Supported Platforms:
-        ``GPU``
-    """
-
-    @prim_attr_register
-    def __init__(self):
-        self.init_prim_io_names(
-            inputs=['input_tensor'],
-            outputs=['output_tensor'])
-
-    def infer_shape(self, input_shape):
-        self.add_prim_attr('fftx', input_shape[0])
-        self.add_prim_attr('ffty', input_shape[1])
-        return [input_shape[0], 2*(input_shape[1]-1)]
-
-    def infer_dtype(self, input_dtype):
-        validator.check_tensor_dtype_valid('input_tensor', input_dtype, [mstype.complex64], self.name)
-        return mstype.float32
-
-
-class FFT3D(PrimitiveWithInfer):
-    """
-    Forward FFT with Three-Dimensional Input.
-
-    .. warning::
-        This is an experimental prototype that is subject to change and/or deletion.
-
-    Inputs:
-        - **input_tensor** (Tensor) - Three dimensional tensor, supported
-          data type is float32.
-
-    Outputs:
-        - **output_tensor** (Tensor) - The tensor after undergoing fast Fourier
-          transform, the data type is complex64.
-
-    Supported Platforms:
-        ``GPU``
-    """
-
-    @prim_attr_register
-    def __init__(self):
-        self.init_prim_io_names(
-            inputs=['input_tensor'],
-            outputs=['output_tensor'])
-
-    def infer_shape(self, input_shape):
-        self.add_prim_attr('fftx', input_shape[0])
-        self.add_prim_attr('ffty', input_shape[1])
-        self.add_prim_attr('fftz', input_shape[2])
-        return [input_shape[0], input_shape[1], int(input_shape[2]/2)+1]
-
-    def infer_dtype(self, input_dtype):
-        validator.check_tensor_dtype_valid('input_tensor', input_dtype, [mstype.float32], self.name)
-        return mstype.complex64
-
-
-class IFFT3D(PrimitiveWithInfer):
-    """
-    Inverse FFT with Three-Dimensional Input.
-
-    .. warning::
-        This is an experimental prototype that is subject to change and/or deletion.
-
-    Inputs:
-        - **input_tensor** (Tensor) - Three dimensional input tensor, supported data
-          type is complex64.
-
-    Outputs:
-        - **output_tensor** (Tensor) - Returns the tensor after undergoing
-          inverse Fourier transform, the data type is float32.
-
-    Supported Platforms:
-        ``GPU``
-    """
-
-    @prim_attr_register
-    def __init__(self):
-        self.init_prim_io_names(
-            inputs=['input_tensor'],
-            outputs=['output_tensor'])
-
-    def infer_shape(self, input_shape):
-        self.add_prim_attr('fftx', input_shape[0])
-        self.add_prim_attr('ffty', input_shape[1])
-        self.add_prim_attr('fftz', input_shape[2])
-        return [input_shape[0], input_shape[1], (input_shape[2]-1)*2]
-
-    def infer_dtype(self, input_dtype):
-        validator.check_tensor_dtype_valid('input_tensor', input_dtype, [mstype.complex64], self.name)
-        return mstype.float32
-
-
-class NeighborListUpdate(PrimitiveWithInfer):
-    """
-    Update (or construct if first time) the Verlet neighbor list for the
-    calculation of short-ranged force. Assume the number of atoms is N,
-    the number of grids divided is G, the maximum number of atoms in one
-    grid is m, the maximum number of atoms in single atom's neighbor list
-    is L, and the number of total atom in excluded list is E.
-
-    Args:
-        grid_numbers (int32): the total number of grids divided G.
-        atom_numbers (int32): the number of atoms n.
-        not_first_time (int32): whether to construct the neighbor list first time or not.
-        nxy (int32): the total number of grids divided in xy plane.
-        excluded_atom_numbers (int32): the total atom numbers in the excluded list E.
-        cutoff_square (float32): the cutoff square distance for short-range force calculation.
-        half_skin_square (float32): the maximum square value of the distance atom allowed to move between two updates.
-        cutoff_with_skin (float32): cutoff + skin, indicates the radius of the neighbor list for each atom.
-        half_cutoff_with_skin (float32): cutoff_with_skin/2.
-        cutoff_with_skin_square (float32): the square value of cutoff_with_skin.
-        refresh_interval (int32): the number of iteration steps between two updates of neighbor list. Default: 20.
-        cutoff (float32): the cutoff distance for short-range force calculation. Default: 10.0.
-        skin (float32): the maximum value of the distance atom allowed to move. Default: 2.0.
-        max_atom_in_grid_numbers (int32): the maximum number of atoms in one grid k. Default: 64.
-        max_neighbor_numbers (int32): The maximum number of neighbors m. Default: 800.
-
-    Inputs:
-        - **atom_numbers_in_grid_bucket** (Tensor) - The number of atoms in each grid bucket.
-          The data type is int32 and the shape is :math:`(G,)`.
-        - **bucket** (Tensor) - (Tensor) - The atom indices in each grid bucket.
-          The data type is int32 and the shape is :math:`(G, k)`.
-        - **crd** (Tensor) - The coordinates of each atom.
-          The data type is float32 and the shape is :math:`(n, 3)`.
-        - **box_length** (Tensor) - The box length of the simulation box.
-          The data type is float32 and the shape is :math:`(3,)`.
-        - **grid_N** (Tensor) - The number of grids divided of 3 dimensions of the simulation box.
-          The data type is int32 and the shape is :math:`(3,)`.
-        - **grid_length_inverse** (Tensor) - The inverse value of grid length.
-          The data type is float32 and the shape is :math:`(3,)`.
-        - **atom_in_grid_serial** (Tensor) - The grid index for each atom.
-          The data type is int32 and the shape is :math:`(n,)`.
-        - **old_crd** (Tensor) - The coordinates before update of each atom.
-          The data type is float32 and the shape is :math:`(n, 3)`.
-        - **crd_to_uint_crd_cof** (Tensor) - The scale factor between the unsigned int coordinate and the real one.
-          The data type is float32 and the shape is :math:`(3,)`.
-        - **uint_crd** (Tensor) - The unsigned int coordinates value fo each atom.
-          The data type is unsigned int32 and the shape is :math:`(n, 3)`.
-        - **gpointer** (Tensor) - The nearest neighbor grids (including self) of each grid.
-          The data type is int32 and the shape is :math:`(G, 125)`.
-        - **nl_atom_numbers** (Tensor) - The number of atoms in neighbor list of each atom.
-          The data type is int32 and the shape is :math:`(n,)`.
-        - **nl_atom_serial** (Tensor) - The indices of atoms in neighbor list of each atom.
-          The data type is int32 and the shape is :math:`(n, m)`.
-        - **uint_dr_to_dr_cof** (Tensor) - The scale factor.
-          The data type is float32 and the shape is :math:`(3,)`.
-        - **excluded_list_start** (Tensor) - The start excluded index in excluded list for each atom.
-          The data type is int32 and the shape is :math:`(n,)`.
-        - **excluded_list** (Tensor) - The contiguous join of excluded list of each atom.
-          The data type is int32 and the shape is :math:`(E,)`.
-        - **excluded_numbers** (Tensor) - The number of atom excluded in excluded list for each atom.
-          The data type is int32 and the shape is :math:`(n,)`.
-        - **need_refresh_flag** (Tensor) - Whether the neighbor list of each atom need update or not.
-          The data type is int32 and the shape is :math:`(1,)`.
-        - **refresh_count** (Tensor) - Count how many iteration steps have passed since last update.
-          The data type is int32 and the shape is :math:`(1,)` or :math:`()`.
-
-    Outputs:
-        - **res** (Tensor) - The return value after updating successfully.
-          The data type is float32 and the shape is :math:`(1,)`.
-
-    Supported Platforms:
-        ``GPU``
-    """
-
-    @prim_attr_register
-    def __init__(self, grid_numbers, atom_numbers, not_first_time, nxy, excluded_atom_numbers,
-                 cutoff_square, half_skin_square, cutoff_with_skin, half_cutoff_with_skin, cutoff_with_skin_square,
-                 refresh_interval=20, cutoff=10.0, skin=2.0, max_atom_in_grid_numbers=64, max_neighbor_numbers=800):
-        self.grid_numbers = grid_numbers
-        self.atom_numbers = atom_numbers
-        self.refresh_interval = refresh_interval
-        self.not_first_time = not_first_time
-        self.cutoff = cutoff
-        self.skin = skin
-        self.max_atom_in_grid_numbers = max_atom_in_grid_numbers
-        self.nxy = nxy
-        self.excluded_atom_numbers = excluded_atom_numbers
-        self.cutoff_square = cutoff_square
-        self.half_skin_square = half_skin_square
-        self.cutoff_with_skin = cutoff_with_skin
-        self.half_cutoff_with_skin = half_cutoff_with_skin
-        self.cutoff_with_skin_square = cutoff_with_skin_square
-        self.max_neighbor_numbers = max_neighbor_numbers
-        validator.check_value_type('grid_numbers', grid_numbers, int, self.name)
-        validator.check_value_type('atom_numbers', atom_numbers, int, self.name)
-        validator.check_value_type('refresh_interval', refresh_interval, int, self.name)
-        validator.check_value_type('not_first_time', not_first_time, int, self.name)
-        validator.check_value_type('cutoff', cutoff, float, self.name)
-        validator.check_value_type('skin', skin, float, self.name)
-        validator.check_value_type('max_atom_in_grid_numbers', max_atom_in_grid_numbers, int, self.name)
-        validator.check_value_type('nxy', nxy, int, self.name)
-        validator.check_value_type('excluded_atom_numbers', excluded_atom_numbers, int, self.name)
-        validator.check_value_type('cutoff_square', cutoff_square, float, self.name)
-        validator.check_value_type('half_skin_square', half_skin_square, float, self.name)
-        validator.check_value_type('cutoff_with_skin', cutoff_with_skin, float, self.name)
-        validator.check_value_type('half_cutoff_with_skin', half_cutoff_with_skin, float, self.name)
-        validator.check_value_type('cutoff_with_skin_square', cutoff_with_skin_square, float, self.name)
-        validator.check_value_type('max_neighbor_numbers', max_neighbor_numbers, int, self.name)
-        self.init_prim_io_names(
-            inputs=['atom_numbers_in_grid_bucket', 'bucket', 'crd', 'box_length', 'grid_N', 'grid_length_inverse',
-                    'atom_in_grid_serial', 'old_crd', 'crd_to_uint_crd_cof', 'uint_crd', 'gpointer', 'nl_atom_numbers',
-                    'nl_atom_serial', 'uint_dr_to_dr_cof', 'excluded_list_start', 'excluded_list', 'excluded_numbers',
-                    'need_refresh_flag', 'refresh_count'], outputs=['res'])
-
-        self.add_prim_attr('grid_numbers', self.grid_numbers)
-        self.add_prim_attr('atom_numbers', self.atom_numbers)
-        self.add_prim_attr('refresh_interval', self.refresh_interval)
-        self.add_prim_attr('not_first_time', self.not_first_time)
-        self.add_prim_attr('cutoff', self.cutoff)
-        self.add_prim_attr('skin', self.skin)
-        self.add_prim_attr('max_atom_in_grid_numbers', self.max_atom_in_grid_numbers)
-        self.add_prim_attr('nxy', self.nxy)
-        self.add_prim_attr('excluded_atom_numbers', self.excluded_atom_numbers)
-        self.add_prim_attr('cutoff_square', self.cutoff_square)
-        self.add_prim_attr('half_skin_square', self.half_skin_square)
-        self.add_prim_attr('cutoff_with_skin', self.cutoff_with_skin)
-        self.add_prim_attr('half_cutoff_with_skin', self.half_cutoff_with_skin)
-        self.add_prim_attr('cutoff_with_skin_square', self.cutoff_with_skin_square)
-        self.add_prim_attr('side_effect_mem', True)
-
-    def infer_shape(self, atom_numbers_in_grid_bucket_shape, bucket_shape, crd_shape, box_length_shape, grid_n_shape,
-                    grid_length_inverse_shape, atom_in_grid_serial_shape, old_crd_shape, crd_to_uint_crd_cof_shape,
-                    uint_crd_shape, gpointer_shape, nl_atom_numbers_shape, nl_atom_serial_shape,
-                    uint_dr_to_dr_cof_shape, excluded_list_start_shape, excluded_list_shape, excluded_numbers_shape,
-                    need_refresh_flag_shape, refresh_count_shape):
-        validator.check_int(len(atom_numbers_in_grid_bucket_shape), 1, Rel.EQ,
-                            "atom_numbers_in_grid_bucket_dim", self.name)
-        validator.check_int(len(bucket_shape), 2, Rel.EQ, "bucket_dim", self.name)
-        validator.check_int(len(crd_shape), 2, Rel.EQ, "crd_dim", self.name)
-        validator.check_int(len(box_length_shape), 1, Rel.EQ, "box_length_dim", self.name)
-        validator.check_int(len(grid_n_shape), 1, Rel.EQ, "grid_n_dim", self.name)
-        validator.check_int(len(grid_length_inverse_shape), 1, Rel.EQ, "grid_length_inverse_dim", self.name)
-        validator.check_int(len(atom_in_grid_serial_shape), 1, Rel.EQ, "atom_in_grid_serial_dim", self.name)
-        validator.check_int(len(old_crd_shape), 2, Rel.EQ, "old_crd_dim", self.name)
-        validator.check_int(len(crd_to_uint_crd_cof_shape), 1, Rel.EQ, "crd_to_uint_crd_cof_dim", self.name)
-        validator.check_int(len(uint_crd_shape), 2, Rel.EQ, "uint_crd_dim", self.name)
-        validator.check_int(len(gpointer_shape), 2, Rel.EQ, "gpointer_dim", self.name)
-        validator.check_int(len(nl_atom_numbers_shape), 1, Rel.EQ, "nl_atom_numbers_dim", self.name)
-        validator.check_int(len(nl_atom_serial_shape), 2, Rel.EQ, "nl_atom_serial_dim", self.name)
-        validator.check_int(len(uint_dr_to_dr_cof_shape), 1, Rel.EQ, "uint_dr_to_dr_cof_dim", self.name)
-        validator.check_int(len(excluded_list_start_shape), 1, Rel.EQ, "excluded_list_start_dim", self.name)
-        validator.check_int(len(excluded_list_shape), 1, Rel.EQ, "excluded_list_dim", self.name)
-        validator.check_int(len(excluded_numbers_shape), 1, Rel.EQ, "excluded_numbers_dim", self.name)
-        validator.check_int(len(need_refresh_flag_shape), 1, Rel.EQ, "need_refresh_flag_dim", self.name)
-        validator.check_int(len(refresh_count_shape), 1, Rel.LE, "refresh_count_dim", self.name)
-        validator.check_int(atom_numbers_in_grid_bucket_shape[0], self.grid_numbers, Rel.EQ,
-                            "atom_numbers_in_grid_bucket", self.name)
-        validator.check_int(bucket_shape[0], self.grid_numbers, Rel.EQ, "bucket", self.name)
-        validator.check_int(bucket_shape[1], self.max_atom_in_grid_numbers, Rel.EQ, "bucket", self.name)
-        validator.check_int(crd_shape[0], self.atom_numbers, Rel.EQ, "crd", self.name)
-        validator.check_int(crd_shape[1], 3, Rel.EQ, "crd", self.name)
-        validator.check_int(box_length_shape[0], 3, Rel.EQ, "box_length", self.name)
-        validator.check_int(grid_n_shape[0], 3, Rel.EQ, "grid_N", self.name)
-        validator.check_int(grid_length_inverse_shape[0], 3, Rel.EQ, "grid_length_inverse", self.name)
-        validator.check_int(atom_in_grid_serial_shape[0], self.atom_numbers, Rel.EQ, "atom_in_grid_serial",
-                            self.name)
-        validator.check_int(old_crd_shape[0], self.atom_numbers, Rel.EQ, "old_crd", self.name)
-        validator.check_int(old_crd_shape[1], 3, Rel.EQ, "old_crd", self.name)
-        validator.check_int(crd_to_uint_crd_cof_shape[0], 3, Rel.EQ, "crd_to_uint_crd_cof", self.name)
-        validator.check_int(uint_crd_shape[0], self.atom_numbers, Rel.EQ, "uint_crd", self.name)
-        validator.check_int(uint_crd_shape[1], 3, Rel.EQ, "uint_crd", self.name)
-        validator.check_int(gpointer_shape[0], self.grid_numbers, Rel.EQ, "gpointer", self.name)
-        validator.check_int(gpointer_shape[1], 125, Rel.EQ, "gpointer", self.name)
-        validator.check_int(nl_atom_numbers_shape[0], self.atom_numbers, Rel.EQ, "nl_atom_numbers", self.name)
-        validator.check_int(nl_atom_serial_shape[0], self.atom_numbers, Rel.EQ, "nl_atom_serial", self.name)
-        validator.check_int(nl_atom_serial_shape[1], self.max_neighbor_numbers, Rel.EQ, "nl_atom_serial",
-                            self.name)
-        validator.check_int(uint_dr_to_dr_cof_shape[0], 3, Rel.EQ, "uint_dr_to_dr_cof", self.name)
-        validator.check_int(excluded_list_start_shape[0], self.atom_numbers, Rel.EQ, "excluded_list_start",
-                            self.name)
-        validator.check_int(excluded_list_shape[0], self.excluded_atom_numbers, Rel.EQ, "excluded_list",
-                            self.name)
-        validator.check_int(excluded_numbers_shape[0], self.atom_numbers, Rel.EQ, "excluded_numbers", self.name)
-        validator.check_int(need_refresh_flag_shape[0], 1, Rel.EQ, "need_refresh_flag", self.name)
-        if refresh_count_shape:
-            validator.check_int(refresh_count_shape[0], 1, Rel.EQ, "refresh_count_shape", self.name)
-        return [1,]
-
-    def infer_dtype(self, atom_numbers_in_grid_bucket_dtype, bucket_dtype, crd_dtype, box_length_dtype, grid_n_dtype,
-                    grid_length_inverse_dtype, atom_in_grid_serial_dtype, old_crd_dtype, crd_to_uint_crd_cof_dtype,
-                    uint_crd_dtype, gpointer_dtype, nl_atom_numbers_dtype, nl_atom_serial_dtype,
-                    uint_dr_to_dr_cof_dtype, excluded_list_start_dtype, excluded_list_dtype, excluded_numbers_dtype,
-                    need_refresh_flag_dtype, refresh_count_dtype):
-        validator.check_tensor_dtype_valid('atom_numbers_in_grid_bucket', atom_numbers_in_grid_bucket_dtype,
-                                           [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('bucket', bucket_dtype, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('crd', crd_dtype, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('box_length', box_length_dtype, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('grid_N', grid_n_dtype, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('grid_length_inverse', grid_length_inverse_dtype, [mstype.float32],
-                                           self.name)
-        validator.check_tensor_dtype_valid('atom_in_grid_serial', atom_in_grid_serial_dtype, [mstype.int32],
-                                           self.name)
-        validator.check_tensor_dtype_valid('old_crd', old_crd_dtype, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('crd_to_uint_crd_cof', crd_to_uint_crd_cof_dtype, [mstype.float32],
-                                           self.name)
-        validator.check_tensor_dtype_valid('uint_crd', uint_crd_dtype, [mstype.uint32], self.name)
-        validator.check_tensor_dtype_valid('gpointer', gpointer_dtype, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('nl_atom_numbers', nl_atom_numbers_dtype, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('nl_atom_serial', nl_atom_serial_dtype, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('uint_dr_to_dr_cof', uint_dr_to_dr_cof_dtype, [mstype.float32],
-                                           self.name)
-        validator.check_tensor_dtype_valid('excluded_list_start', excluded_list_start_dtype, [mstype.int32],
-                                           self.name)
-        validator.check_tensor_dtype_valid('excluded_list', excluded_list_dtype, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('excluded_numbers', excluded_numbers_dtype, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('need_refresh_flag', need_refresh_flag_dtype, [mstype.int32],
-                                           self.name)
-        validator.check_tensor_dtype_valid('refresh_count', refresh_count_dtype, [mstype.int32],
-                                           self.name)
-        return mstype.float32