mindspore 1.10.0__cp39-cp39-win_amd64.whl → 2.0.0rc1__cp39-cp39-win_amd64.whl

mindspore/ops/operations/sponge_update_ops.py

@@ -1,2546 +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 for new modules. 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 RefreshUintCrd(PrimitiveWithInfer):
-    """
-    Refresh the unsigned coordinate of each constrained atom in each constrain iteration.
-
-    .. warning::
-        This is an experimental prototype that is subject to change and/or deletion.
-
-    Args:
-        atom_numbers (int32): the number of atoms n.
-        half_exp_gamma_plus_half (float32): constant value (1.0 + exp(gamma * dt)) if Langvin-Liu thermostat is used,
-            where gamma is friction coefficient and dt is the simulation time step, 1.0 otherwise.
-
-    Inputs:
-        - **crd** (Tensor) - The coordinate of each atom.
-          The data type is float32 and the shape is :math:`(n, 3)`.
-        - **quarter_cof** (Tensor) - The 3-D scale factor.
-          The data type is float32 and the shape is :math:`(3,)`.
-        - **test_frc** (Tensor) - The constraint force.
-          The data type is float32 and the shape is :math:`(n, 3)`.
-        - **mass_inverse** (Tensor) - The inverse value of mass of each atom.
-          The data type is float32 and the shape is :math:`(n,)`.
-
-    Outputs:
-        - **uint_crd** (Tensor) - The unsigned int coordinate value of each atom.
-          The data type is uint32 and the shape is :math:`(n, 3)`.
-
-    Supported Platforms:
-        ``GPU``
-    """
-
-    @prim_attr_register
-    def __init__(self, atom_numbers, half_exp_gamma_plus_half):
-        validator.check_value_type('atom_numbers', atom_numbers, int, self.name)
-        validator.check_value_type('half_exp_gamma_plus_half', half_exp_gamma_plus_half, float, self.name)
-        self.atom_numbers = atom_numbers
-        self.half_exp_gamma_plus_half = half_exp_gamma_plus_half
-        self.add_prim_attr('atom_numbers', self.atom_numbers)
-        self.add_prim_attr('half_exp_gamma_plus_half', self.half_exp_gamma_plus_half)
-        self.init_prim_io_names(
-            inputs=['crd', 'quarter_cof', 'test_frc', 'mass_inverse'],
-            outputs=['uint_crd'])
-
-    def infer_shape(self, crd_shape, quarter_cof_shape, test_frc_shape, mass_inverse_shape):
-        cls_name = self.name
-        n = self.atom_numbers
-        validator.check_int(len(crd_shape), 2, Rel.EQ, "crd_dim", cls_name)
-        validator.check_int(len(quarter_cof_shape), 1, Rel.EQ, "quarter_cof_dim", cls_name)
-        validator.check_int(len(test_frc_shape), 2, Rel.EQ, "test_frc_dim", cls_name)
-        validator.check_int(len(mass_inverse_shape), 1, Rel.EQ, "mass_inverse_dim", cls_name)
-
-        validator.check_int(crd_shape[0], n, Rel.EQ, "crd_shape[0]", cls_name)
-        validator.check_int(crd_shape[1], 3, Rel.EQ, "crd_shape[1]", cls_name)
-        validator.check_int(quarter_cof_shape[0], 3, Rel.EQ, "quarter_cof_shape", cls_name)
-        validator.check_int(test_frc_shape[0], n, Rel.EQ, "test_frc_shape[0]", cls_name)
-        validator.check_int(test_frc_shape[1], 3, Rel.EQ, "test_frc_shape[1]", cls_name)
-        validator.check_int(mass_inverse_shape[0], n, Rel.EQ, "mass_inverse_shape", cls_name)
-
-        return [n, 3]
-
-    def infer_dtype(self, crd_dtype, quarter_cof_dtype, test_frc_dtype, mass_inverse_dtype):
-        validator.check_tensor_dtype_valid('crd', crd_dtype, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('quarter_cof', quarter_cof_dtype, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('test_frc', test_frc_dtype, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('mass_inverse', mass_inverse_dtype, [mstype.float32], self.name)
-        return mstype.uint32
-
-
-class ConstrainForceCycleWithVirial(PrimitiveWithInfer):
-    """
-    Calculate the constraint force and virial in each iteration.
-
-    .. warning::
-        This is an experimental prototype that is subject to change and/or deletion.
-
-    Args:
-        atom_numbers (int32): the number of atoms n.
-        constrain_pair_numbers (int32): the number of constrain pairs m.
-
-    Inputs:
-        - **uint_crd** (Tensor) - The unsigned int coordinate value of each atom.
-          The data type is uint32 and the shape is :math:`(n, 3)`.
-        - **scaler** (Tensor) - The 3-D scale factor (x, y, z),
-          The data type is float32 and the shape is :math:`(3,)`.
-        - **pair_dr** (Tensor) - The displacement vector of each constrained atom pair.
-          The data type is float32 and the shape is :math:`(m, 3)`.
-        - **atom_i_serials** (Tensor) - The first atom index of each constrained atom pair.
-          The data type is int32 and the shape is :math:`(m,)`.
-        - **atom_j_serials** (Tensor) - The second atom index of each constrained atom pair.
-          The data type is int32 and the shape is :math:`(m,)`.
-        - **constant_rs** (Tensor) - The constrained distance of each constrained atom pair.
-          The data type is float32 and the shape is :math:`(m,)`.
-        - **constrain_ks** (Tensor) - The coefficient of each constrained atom pair.
-          The data type is float32 and the shape is :math:`(m,)`.
-
-    Outputs:
-        - **test_frc** (Tensor) - The constraint force.
-          The data type is float32 and the shape is :math:`(n, 3)`.
-        - **atom_virial** (Tensor) - The virial caused by constraint force of each atom.
-          The data type is float32 and the shape is :math:`(m,)`.
-
-    Supported Platforms:
-        ``GPU``
-    """
-
-    @prim_attr_register
-    def __init__(self, atom_numbers, constrain_pair_numbers):
-        validator.check_value_type('atom_numbers', atom_numbers, int, self.name)
-        validator.check_value_type('constrain_pair_numbers', constrain_pair_numbers, int, self.name)
-        self.atom_numbers = atom_numbers
-        self.constrain_pair_numbers = constrain_pair_numbers
-        self.add_prim_attr('atom_numbers', self.atom_numbers)
-        self.add_prim_attr('constrain_pair_numbers', self.constrain_pair_numbers)
-        self.init_prim_io_names(
-            inputs=['uint_crd', 'scaler', 'pair_dr', 'atom_i_serials', 'atom_j_serials',
-                    'constant_rs', 'constrain_ks'],
-            outputs=['test_frc', 'atom_virial'])
-
-    def infer_shape(self, uint_crd_shape, scaler_shape, pair_dr_shape, atom_i_serials_shape,
-                    atom_j_serials_shape, constant_rs_shape, constrain_ks_shape):
-        cls_name = self.name
-        n = self.atom_numbers
-        m = self.constrain_pair_numbers
-        validator.check_int(len(uint_crd_shape), 2, Rel.EQ, "uint_crd_dim", cls_name)
-        validator.check_int(len(scaler_shape), 1, Rel.EQ, "scaler_dim", cls_name)
-        validator.check_int(len(pair_dr_shape), 2, Rel.EQ, "pair_dr_dim", cls_name)
-        validator.check_int(len(atom_i_serials_shape), 1, Rel.EQ, "atom_i_serials_dim", cls_name)
-        validator.check_int(len(atom_j_serials_shape), 1, Rel.EQ, "atom_j_serials_dim", cls_name)
-        validator.check_int(len(constant_rs_shape), 1, Rel.EQ, "constant_rs_dim", cls_name)
-        validator.check_int(len(constrain_ks_shape), 1, Rel.EQ, "constrain_ks_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(scaler_shape[0], 3, Rel.EQ, "scaler_shape", cls_name)
-        validator.check_int(pair_dr_shape[0], m, Rel.EQ, "pair_dr_shape[0]", cls_name)
-        validator.check_int(pair_dr_shape[1], 3, Rel.EQ, "pair_dr_shape[1]", cls_name)
-        validator.check_int(atom_i_serials_shape[0], m, Rel.EQ, "atom_i_serials_shape", cls_name)
-        validator.check_int(atom_j_serials_shape[0], m, Rel.EQ, "atom_j_serials_shape", cls_name)
-        validator.check_int(constant_rs_shape[0], m, Rel.EQ, "constant_rs_shape", cls_name)
-        validator.check_int(constrain_ks_shape[0], m, Rel.EQ, "constrain_ks_shape", cls_name)
-        return [n, 3], [m,]
-
-    def infer_dtype(self, uint_crd_dtype, scaler_dtype, pair_dr_dtype, atom_i_serials_dtype,
-                    atom_j_serials_dtype, constant_rs_dtype, constrain_ks_dtype):
-        validator.check_tensor_dtype_valid('uint_crd', uint_crd_dtype, [mstype.uint32], self.name)
-        validator.check_tensor_dtype_valid('scaler', scaler_dtype, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('pair_dr', pair_dr_dtype, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('atom_i_serials', atom_i_serials_dtype, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('atom_j_serials', atom_j_serials_dtype, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('constant_rs', constant_rs_dtype, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('constrain_ks', constrain_ks_dtype, [mstype.float32], self.name)
-        return mstype.float32, mstype.float32
-
-
-class LastCrdToDr(PrimitiveWithInfer):
-    """
-    Calculate the displacement vector of each constrained atom pair.
-
-    .. warning::
-        This is an experimental prototype that is subject to change and/or deletion.
-
-    Args:
-        atom_numbers (int32): the number of atoms n.
-        constrain_pair_numbers (int32): the number of constrain pairs m.
-
-    Inputs:
-        - **crd** (Tensor) - The coordinate of each atom.
-          The data type is float32 and the shape is :math:`(n, 3)`.
-        - **quarter_cof** (Tensor) - The 3-D scale factor.
-          The data type is float32 and the shape is :math:`(3,)`.
-        - **uint_dr_to_dr** (Tensor) - The 3-D scale factor (x, y, z)
-          The data type is int32 and the shape is :math:`(3,)`..
-        - **atom_i_serials** (Tensor) - The first atom index of each constrained atom pair.
-          The data type is int32 and the shape is :math:`(m,)`.
-        - **atom_j_serials** (Tensor) - The second atom index of each constrained atom pair.
-          The data type is int32 and the shape is :math:`(m,)`.
-        - **constant_rs** (Tensor) - The constrained distance of each constrained atom pair.
-          The data type is float32 and the shape is :math:`(m,)`.
-        - **constrain_ks** (Tensor) - The coefficient of each constrained atom pair.
-          The data type is float32 and the shape is :math:`(m,)`.
-
-    Outputs:
-        - **pair_dr** (Tensor) - The displacement vector of each constrained atom pair.
-          The data type is float32 and the shape is :math:`(m, 3)`.
-
-    Supported Platforms:
-        ``GPU``
-    """
-
-    @prim_attr_register
-    def __init__(self, atom_numbers, constrain_pair_numbers):
-        validator.check_value_type('atom_numbers', atom_numbers, int, self.name)
-        validator.check_value_type('constrain_pair_numbers', constrain_pair_numbers, int, self.name)
-        self.constrain_pair_numbers = constrain_pair_numbers
-        self.atom_numbers = atom_numbers
-        self.add_prim_attr('constrain_pair_numbers', self.constrain_pair_numbers)
-        self.init_prim_io_names(
-            inputs=['crd', 'quarter_cof', 'uint_dr_to_dr', 'atom_i_serials', 'atom_j_serials',
-                    'constant_rs', 'constrain_ks'],
-            outputs=['pair_dr'])
-
-    def infer_shape(self, crd_shape, quarter_cof_shape, uint_dr_to_dr_shape, atom_i_serials_shape,
-                    atom_j_serials_shape, constant_rs_shape, constrain_ks_shape):
-        cls_name = self.name
-        n = self.atom_numbers
-        m = self.constrain_pair_numbers
-        validator.check_int(len(crd_shape), 2, Rel.EQ, "crd_dim", cls_name)
-        validator.check_int(len(quarter_cof_shape), 1, Rel.EQ, "quarter_cof_dim", cls_name)
-        validator.check_int(len(uint_dr_to_dr_shape), 1, Rel.EQ, "quarter_cof_dim", cls_name)
-        validator.check_int(len(atom_i_serials_shape), 1, Rel.EQ, "atom_i_serials_dim", cls_name)
-        validator.check_int(len(atom_j_serials_shape), 1, Rel.EQ, "atom_j_serials_dim", cls_name)
-        validator.check_int(len(constant_rs_shape), 1, Rel.EQ, "constant_rs_dim", cls_name)
-        validator.check_int(len(constrain_ks_shape), 1, Rel.EQ, "constrain_ks_dim", cls_name)
-
-        validator.check_int(crd_shape[0], n, Rel.EQ, "crd_shape[0]", cls_name)
-        validator.check_int(crd_shape[1], 3, Rel.EQ, "crd_shape[1]", cls_name)
-        validator.check_int(quarter_cof_shape[0], 3, Rel.EQ, "quarter_cof_shape", cls_name)
-        validator.check_int(uint_dr_to_dr_shape[0], 3, Rel.EQ, "uint_dr_to_dr_shape", cls_name)
-        validator.check_int(atom_i_serials_shape[0], m, Rel.EQ, "atom_i_serials_shape", cls_name)
-        validator.check_int(atom_j_serials_shape[0], m, Rel.EQ, "atom_j_serials_shape", cls_name)
-        validator.check_int(constant_rs_shape[0], m, Rel.EQ, "constant_rs_shape", cls_name)
-        validator.check_int(constrain_ks_shape[0], m, Rel.EQ, "constrain_ks_shape", cls_name)
-
-        return [m, 3]
-
-    def infer_dtype(self, crd_dtype, quarter_cof_dtype, uint_dr_to_dr_dtype, atom_i_serials_dtype,
-                    atom_j_serials_dtype, constant_rs_dtype, constrain_ks_dtype):
-        validator.check_tensor_dtype_valid('crd', crd_dtype, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('quarter_cof', quarter_cof_dtype, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('uint_dr_to_dr', uint_dr_to_dr_dtype, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('atom_i_serials', atom_i_serials_dtype, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('atom_j_serials', atom_j_serials_dtype, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('constant_rs', constant_rs_dtype, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('constrain_ks', constrain_ks_dtype, [mstype.float32], self.name)
-        return mstype.float32
-
-
-class RefreshCrdVel(PrimitiveWithInfer):
-    """
-    Refresh the coordinate and velocity of each constrained atom after all iterations have ended.
-
-    .. warning::
-        This is an experimental prototype that is subject to change and/or deletion.
-
-    Args:
-        atom_numbers (int32): the number of atoms n.
-        dt_inverse (float32): the inverse value of simulation time step.
-        dt (float32): the simulation time step.
-        exp_gamma (float32): constant value exp(gamma * dt).
-        half_exp_gamma_plus_half (float32): constant value (1 + exp_gamma)/2.
-
-    Inputs:
-        - **crd** (Tensor) - The coordinate of each atom.
-          The data type is float32 and the shape is :math:`(n, 3)`.
-        - **vel** (Tensor) - The velocity of each atom.
-          The data type is float32 and the shape is :math:`(n, 3)`.
-        - **test_frc** (Tensor) - The constraint force calculated in the last iteration.
-          The data type is float32 and the shape is :math:`(n, 3)`.
-        - **mass_inverse** (Tensor) - The inverse value of mass of each atom.
-          The data type is float32 and the shape is :math:`(n,)`.
-
-    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, atom_numbers, dt_inverse, dt, exp_gamma, half_exp_gamma_plus_half):
-        validator.check_value_type('atom_numbers', atom_numbers, int, self.name)
-        validator.check_value_type('dt', dt, float, self.name)
-        validator.check_value_type('dt_inverse', dt_inverse, float, self.name)
-        validator.check_value_type('exp_gamma', exp_gamma, float, self.name)
-        validator.check_value_type('half_exp_gamma_plus_half', half_exp_gamma_plus_half, float, self.name)
-        self.atom_numbers = atom_numbers
-        self.dt_inverse = dt_inverse
-        self.dt = dt
-        self.exp_gamma = exp_gamma
-        self.half_exp_gamma_plus_half = half_exp_gamma_plus_half
-        self.add_prim_attr('atom_numbers', self.atom_numbers)
-        self.add_prim_attr('dt_inverse', self.dt_inverse)
-        self.add_prim_attr('dt', self.dt)
-        self.add_prim_attr('exp_gamma', self.exp_gamma)
-        self.add_prim_attr('half_exp_gamma_plus_half', self.half_exp_gamma_plus_half)
-        self.init_prim_io_names(
-            inputs=['crd', 'vel', 'test_frc', 'mass_inverse'],
-            outputs=['res'])
-        self.add_prim_attr('side_effect_mem', True)
-
-    def infer_shape(self, crd_shape, vel_shape, test_frc_shape, mass_inverse_shape):
-        cls_name = self.name
-        n = self.atom_numbers
-        validator.check_int(len(crd_shape), 2, Rel.EQ, "crd_dim", cls_name)
-        validator.check_int(len(vel_shape), 2, Rel.EQ, "vel_dim", cls_name)
-        validator.check_int(len(test_frc_shape), 2, Rel.EQ, "test_frc_dim", cls_name)
-        validator.check_int(len(mass_inverse_shape), 1, Rel.EQ, "mass_inverse_dim", cls_name)
-
-        validator.check_int(crd_shape[0], n, Rel.EQ, "crd_shape[0]", cls_name)
-        validator.check_int(crd_shape[1], 3, Rel.EQ, "crd_shape[1]", cls_name)
-        validator.check_int(vel_shape[0], n, Rel.EQ, "vel_shape[0]", cls_name)
-        validator.check_int(vel_shape[1], 3, Rel.EQ, "vel_shape[1]", cls_name)
-        validator.check_int(test_frc_shape[0], n, Rel.EQ, "test_frc_shape[0]", cls_name)
-        validator.check_int(test_frc_shape[1], 3, Rel.EQ, "test_frc_shape[1]", cls_name)
-        validator.check_int(mass_inverse_shape[0], n, Rel.EQ, "mass_inverse_shape[0]", cls_name)
-        return [1,]
-
-    def infer_dtype(self, crd_dtype, vel_dtype, test_frc_dtype, mass_inverse_dtype):
-        validator.check_tensor_dtype_valid('crd', crd_dtype, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('vel', vel_dtype, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('test_frc', test_frc_dtype, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('mass_inverse', mass_inverse_dtype, [mstype.float32], self.name)
-        return mstype.float32
-
-
-class CalculateNowrapCrd(PrimitiveWithInfer):
-    """
-    Calculate the inside-box periodic image of each atom.
-
-    .. warning::
-        This is an experimental prototype that is subject to change and/or deletion.
-
-    Args:
-        atom_numbers (int32): the number of atoms n.
-
-    Inputs:
-        - **crd** (Tensor) - The coordinate of each atom.
-          The data type is float32 and the shape is :math:`(n, 3)`.
-        - **box** (Tensor) - The 3-D size of system.
-          The data type is float32 and the shape is :math:`(3, )`.
-        - **box_map_times** (Tensor) - The number of times each atom has crossed the box.
-          The data type is int32 and the shape is :math:`(n, 3)`.
-
-    Outputs:
-        - **nowrap_crd** (Tensor) - The inside-box periodic image of 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):
-        validator.check_value_type('atom_numbers', atom_numbers, int, self.name)
-        self.atom_numbers = atom_numbers
-        self.add_prim_attr('atom_numbers', self.atom_numbers)
-        self.init_prim_io_names(
-            inputs=['crd', 'box', 'box_map_times'],
-            outputs=['nowrap_crd'])
-
-    def infer_shape(self, crd_shape, box_shape, box_map_times_shape):
-        cls_name = self.name
-        n = self.atom_numbers
-        validator.check_int(len(crd_shape), 2, Rel.EQ, "crd_dim", cls_name)
-        validator.check_int(len(box_shape), 1, Rel.EQ, "box_dim", cls_name)
-        validator.check_int(len(box_map_times_shape), 2, Rel.EQ, "box_map_times_dim", cls_name)
-
-        validator.check_int(crd_shape[0], n, Rel.EQ, "crd_shape[0]", cls_name)
-        validator.check_int(crd_shape[1], 3, Rel.EQ, "crd_shape[1]", cls_name)
-        validator.check_int(box_shape[0], 3, Rel.EQ, "box_shape[0]", cls_name)
-        validator.check_int(box_map_times_shape[0], n, Rel.EQ, "box_map_times_shape[0]", cls_name)
-        validator.check_int(box_map_times_shape[1], 3, Rel.EQ, "box_map_times_shape[1]", cls_name)
-        return [n, 3]
-
-    def infer_dtype(self, crd_dtype, box_dtype, box_map_times_dtype):
-        validator.check_tensor_dtype_valid('crd', crd_dtype, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('box', box_dtype, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('box_map_times', box_map_times_dtype,
-                                           [mstype.float32], self.name)
-        return mstype.float32
-
-
-class RefreshBoxmapTimes(PrimitiveWithInfer):
-    """
-    Refresh the box-crossing times of each atom.
-
-    .. warning::
-        This is an experimental prototype that is subject to change and/or deletion.
-
-    Args:
-        atom_numbers (int32): the number of atoms n.
-
-    Inputs:
-        - **crd** (Tensor) - The coordinate of each atom.
-          The data type is float32 and the shape is :math:`(n, 3)`.
-        - **old_crd** (Tensor) - The coordinate of each atom at last update.
-          The data type is float32 and the shape is :math:`(n, 3)`.
-        - **box_length_inverse** (Tensor) - The inverse value of box length in 3 dimensions.
-          The data type is float32 and the shape is :math:`(3,)`.
-        - **box_map_times** (Tensor) - The number of times each atom has crossed the box.
-          The data type is int32 and the shape is :math:`(n, 3)`.
-
-    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, atom_numbers):
-        validator.check_value_type('atom_numbers', atom_numbers, int, self.name)
-        self.atom_numbers = atom_numbers
-        self.add_prim_attr('atom_numbers', self.atom_numbers)
-        self.init_prim_io_names(
-            inputs=['crd', 'old_crd', 'box_length_inverse', 'box_map_times'],
-            outputs=['res'])
-
-    def infer_shape(self, crd_shape, old_crd_shape, box_length_inverse_shape, box_map_times_shape):
-        cls_name = self.name
-        n = self.atom_numbers
-        validator.check_int(len(crd_shape), 2, Rel.EQ, "crd_dim", cls_name)
-        validator.check_int(len(old_crd_shape), 2, Rel.EQ, "old_crd_dim", cls_name)
-        validator.check_int(len(box_length_inverse_shape), 1, Rel.EQ, "box_length_inverse_dim", cls_name)
-        validator.check_int(len(box_map_times_shape), 2, Rel.EQ, "box_map_times_dim", cls_name)
-
-        validator.check_int(crd_shape[0], n, Rel.EQ, "crd_shape[0]", cls_name)
-        validator.check_int(crd_shape[1], 3, Rel.EQ, "crd_shape[1]", cls_name)
-        validator.check_int(old_crd_shape[0], n, Rel.EQ, "old_crd_shape[0]", cls_name)
-        validator.check_int(old_crd_shape[1], 3, Rel.EQ, "old_crd_shape[1]", cls_name)
-        validator.check_int(box_length_inverse_shape[0], 3, Rel.EQ, "box_length_inverse_shape[0]", cls_name)
-        validator.check_int(box_map_times_shape[0], n, Rel.EQ, "box_map_times_shape[0]", cls_name)
-        validator.check_int(box_map_times_shape[1], 3, Rel.EQ, "box_map_times_shape[1]", cls_name)
-        return [1,]
-
-    def infer_dtype(self, crd_dtype, old_crd_dtype, box_length_inverse_dtype, box_map_times_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_length_inverse', box_length_inverse_dtype, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('box_map_times', box_map_times_dtype,
-                                           [mstype.int32], self.name)
-        return mstype.float32
-
-
-class Totalc6get(PrimitiveWithInfer):
-    """
-    Get the average dispersion constant of short range Lennard-Jones interaction,
-    for the subsequent long range correction energy and virial. Assume system has m Lennard-Jones types of atoms.
-
-    .. warning::
-        This is an experimental prototype that is subject to change and/or deletion.
-
-    Args:
-        atom_numbers (int32): the number of atoms n.
-
-    Inputs:
-        - **atom_lj_type** (Tensor) - The Lennard-Jones type of each atom.
-          The data type is float32 and the shape is :math:`(n,)`.
-        - **lj_b** (Tensor) - The attraction coefficient of each type. the number of pair atoms is m.
-          The data type is float32 and the shape is :math:`(m,)`.
-
-    Outputs:
-        - **factor** (Tensor) - The average dispersion constant of Lennard-Jones interaction.
-          The data type is float32 and the shape is :math:`(1,)`.
-
-    Supported Platforms:
-        ``GPU``
-    """
-
-    @prim_attr_register
-    def __init__(self, atom_numbers):
-        validator.check_value_type('atom_numbers', atom_numbers, int, self.name)
-        self.atom_numbers = atom_numbers
-        self.add_prim_attr('atom_numbers', self.atom_numbers)
-        self.init_prim_io_names(
-            inputs=['atom_lj_type', 'lj_b'],
-            outputs=['factor'])
-
-    def infer_shape(self, atom_lj_type, lj_b):
-        cls_name = self.name
-        n = self.atom_numbers
-        validator.check_int(len(atom_lj_type), 1, Rel.EQ, "atom_lj_type_dim", cls_name)
-        validator.check_int(len(lj_b), 1, Rel.EQ, "LJ_b_dim", cls_name)
-        validator.check_int(atom_lj_type[0], n, Rel.EQ, "atom_lj_type_shape[0]", cls_name)
-        return [1,]
-
-    def infer_dtype(self, atom_lj_type, lj_b):
-        validator.check_tensor_dtype_valid('atom_lj_type', atom_lj_type, mstype.int32, self.name)
-        validator.check_tensor_dtype_valid('lj_b', lj_b, mstype.float32, self.name)
-        return mstype.float32
-
-
-class CrdToUintCrdQuarter(PrimitiveWithInfer):
-    """
-    Convert FP32 coordinate to Uint32 coordinate.
-
-    .. warning::
-        This is an experimental prototype that is subject to change and/or deletion.
-
-    Args:
-        atom_numbers (int32): the number of atoms n.
-
-    Inputs:
-        - **crd_to_uint_crd_cof** (Tensor) - The crd_to_uint_crd coefficient.
-          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** (Tensor) - The unsigned int coordinates.
-          The data type is unsigned int32 and the shape is :math:`(n, 3)`.
-
-    Supported Platforms:
-        ``GPU``
-    """
-
-    @prim_attr_register
-    def __init__(self, atom_numbers):
-        """Initialize CrdToUintCrdQuarter"""
-        validator.check_value_type('atom_numbers', atom_numbers, int, self.name)
-        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):
-        cls_name = self.name
-        n = self.atom_numbers
-        validator.check_int(len(crd), 2, Rel.EQ, "crd_dim", cls_name)
-        validator.check_int(len(crd_to_uint_crd_cof), 1, Rel.EQ, "crd_to_uint_crd_cof_dim", cls_name)
-        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], n, 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 MDIterationLeapFrogLiujianWithMaxVel(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, but with an maximum velocity limit. 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.
-
-    .. warning::
-        This is an experimental prototype that is subject to change and/or deletion.
-
-    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, exp(-gamma_ln * dt).
-        max_vel (float32): the maximum velocity limit.
-
-    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 sqrt of the 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** (float32) - The output coordinate of 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, half_dt, dt, exp_gamma, max_vel):
-        """Initialize MDIterationLeapFrogLiujian"""
-        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('max_vel', max_vel, float, self.name)
-        self.atom_numbers = atom_numbers
-        self.half_dt = half_dt
-        self.dt = dt
-        self.exp_gamma = exp_gamma
-        self.max_vel = max_vel
-
-        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('max_vel', self.max_vel)
-        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_dim", self.name)
-        validator.check_int(len(sqrt_mass_inverse), 1, Rel.EQ, "sqrt_mass_inverse_dim", 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(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(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 GetCenterOfMass(PrimitiveWithInfer):
-    """
-    Get coordinate of centroid of each residue. Assume system has n atoms.
-
-    .. warning::
-        This is an experimental prototype that is subject to change and/or deletion.
-
-    Args:
-        residue_numbers (int32): the number of residues m.
-
-    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,)`.
-        - **crd** (Tensor) - The coordinate of each atom.
-          The data type is float32 and the shape is :math:`(n, 3)`.
-        - **atom_mass** (Tensor) - The mass of each atom and the atom number is n.
-          The data type is float32 and the shape is :math:`(n,)`.
-        - **residue_mass_inverse** (Tensor) - The inverse of mass of each residue.
-          The data type is float32 and the shape is :math:`(m,)`.
-
-    Outputs:
-        - **center_of_mass** (Tensor) - The coordinate of centroid of each residue.
-          The data type is float32 and the shape is :math:`(m, 3)`.
-
-    Supported Platforms:
-        ``GPU``
-    """
-
-    @prim_attr_register
-    def __init__(self, residue_numbers):
-        """Initialize GetCenterOfMass"""
-        validator.check_value_type('residue_numbers', residue_numbers, int, self.name)
-        self.residue_numbers = residue_numbers
-        self.add_prim_attr('residue_numbers', self.residue_numbers)
-        self.init_prim_io_names(
-            inputs=['start', 'end', 'crd', 'atom_mass', 'residue_mass_inverse'],
-            outputs=['center_of_mass'])
-
-    def infer_shape(self, start, end, crd, atom_mass, residue_mass_inverse):
-        n = crd[0]
-        m = self.residue_numbers
-        validator.check_int(len(start), 1, Rel.EQ, "start_dim", self.name)
-        validator.check_int(len(end), 1, Rel.EQ, "end_dim", self.name)
-        validator.check_int(len(crd), 2, Rel.EQ, "crd_dim", self.name)
-        validator.check_int(len(atom_mass), 1, Rel.EQ, "atom_mass_dim", self.name)
-        validator.check_int(len(residue_mass_inverse), 1, Rel.EQ, "residue_mass_inverse_dim", self.name)
-        validator.check_int(start[0], m, Rel.EQ, "start_shape", self.name)
-        validator.check_int(end[0], m, Rel.EQ, "end_shape", 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(atom_mass[0], n, Rel.EQ, "atom_mass_shape[0]", self.name)
-        validator.check_int(residue_mass_inverse[0], m, Rel.EQ, "residue_mass_inverse_shape", self.name)
-        return [m, 3]
-
-    def infer_dtype(self, start, end, crd, atom_mass, residue_mass_inverse):
-        validator.check_tensor_dtype_valid('start', start, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('end', end, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('crd', crd, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('atom_mass', atom_mass, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('residue_mass_inverse', residue_mass_inverse, [mstype.float32], self.name)
-        return mstype.float32
-
-
-class MapCenterOfMass(PrimitiveWithInfer):
-    """
-    Map all atoms in the same residue to the same periodic box, scale if necessary (usually in pressurestat).
-    Assume system has n atoms.
-
-    .. warning::
-        This is an experimental prototype that is subject to change and/or deletion.
-
-    Args:
-        residue_numbers (int32): the number of residues m.
-
-    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,)`.
-        - **center_of_mass** (Tensor) - The coordinate of centroid of each residue.
-          The data type is float32 and the shape is :math:`(m, 3)`.
-        - **box_length** (Tensor) - The box length of the simulation box.
-          The data type is float32 and the shape is :math:`(3,)`.
-        - **no_wrap_crd** (Tensor) - The coordinate of each atom before wrap.
-          The data type is float32 and the shape is :math:`(n, 3)`.
-        - **crd** (Tensor) - The coordinate of each atom after wrap.
-          The data type is float32 and the shape is :math:`(n, 3)`.
-        - **scaler** (Tensor) - The scaler of system.
-          The data type is float32 and the shape is :math:`(1,)`.
-
-    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, residue_numbers):
-        """Initialize MapCenterOfMass"""
-        validator.check_value_type('residue_numbers', residue_numbers, int, self.name)
-        self.residue_numbers = residue_numbers
-        self.add_prim_attr('residue_numbers', self.residue_numbers)
-        self.init_prim_io_names(
-            inputs=['start', 'end', 'center_of_mass', 'box_length',
-                    'no_wrap_crd', 'crd', 'scaler'],
-            outputs=['res'])
-
-    def infer_shape(self, start, end, center_of_mass, box_length, no_wrap_crd, crd, scaler):
-        m = self.residue_numbers
-        n = crd[0]
-        validator.check_int(len(start), 1, Rel.EQ, "start_dim", self.name)
-        validator.check_int(len(end), 1, Rel.EQ, "end_dim", self.name)
-        validator.check_int(len(center_of_mass), 2, Rel.EQ, "center_of_mass_dim", self.name)
-        validator.check_int(len(box_length), 1, Rel.EQ, "box_length_dim", self.name)
-        validator.check_int(len(no_wrap_crd), 2, Rel.EQ, "no_wrap_crd_dim", self.name)
-        validator.check_int(len(crd), 2, Rel.EQ, "crd_dim", self.name)
-        validator.check_int(len(scaler), 1, Rel.EQ, "scaler_dim", self.name)
-
-        validator.check_int(start[0], m, Rel.EQ, "start_shape", self.name)
-        validator.check_int(end[0], m, Rel.EQ, "end_shape", self.name)
-        validator.check_int(center_of_mass[0], m, Rel.EQ, "center_of_mass_shape[0]", self.name)
-        validator.check_int(center_of_mass[1], 3, Rel.EQ, "center_of_mass_shape[1]", self.name)
-        validator.check_int(box_length[0], 3, Rel.EQ, "box_length_shape", self.name)
-        validator.check_int(scaler[0], 1, Rel.EQ, "scaler_shape", self.name)
-        validator.check_int(no_wrap_crd[0], n, Rel.EQ, "no_wrap_crd_shape[0]", self.name)
-        validator.check_int(no_wrap_crd[1], 3, Rel.EQ, "no_wrap_crd_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)
-        return [1,]
-
-    def infer_dtype(self, start, end, center_of_mass, box_length, no_wrap_crd, crd, scaler):
-        validator.check_tensor_dtype_valid('start', start, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('end', end, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('crd', crd, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('center_of_mass', center_of_mass, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('box_length', box_length, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('no_wrap_crd', no_wrap_crd, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('scaler', scaler, [mstype.float32], self.name)
-        return mstype.float32
-
-
-class NeighborListRefresh(PrimitiveWithInfer):
-    """
-    Update (or construct if first time) the Verlet neighbor list for the
-    calculation of short-ranged force.
-
-    .. warning::
-        This is an experimental prototype that is subject to change and/or deletion.
-
-    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.
-        forced_update (int32): the flag that decides whether to force an update. Default: 0.
-        forced_check (int32): the flag that decides whether to force an check. Default: 0.
-
-    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** (Union[Tensor, Scalar]) - 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,
-                 forced_update=0, forced_check=0):
-        """Initialize NeighborListRefresh"""
-        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
-        self.forced_update = forced_update
-        self.forced_check = forced_check
-        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)
-        validator.check_value_type('forced_update', forced_update, int, self.name)
-        validator.check_value_type('forced_check', forced_check, 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('forced_update', self.forced_update)
-        self.add_prim_attr('forced_check', self.forced_check)
-        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
-
-
-class MDIterationLeapFrog(PrimitiveWithInfer):
-    """
-    One step of classical leap frog algorithm to solve the finite difference
-    Hamiltonian equations of motion for certain system.
-
-    .. warning::
-        This is an experimental prototype that is subject to change and/or deletion.
-
-    Args:
-        atom_numbers (int32): the number of atoms n.
-        dt (float32): the simulation time step.
-
-    Inputs:
-        - **sqrt_mass_inverse** (Tensor) - The square root of the inverse value of the mass 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)`.
-        - **inverse_mass** (Tensor) - The inverse value of mass of each atom.
-          The data type is float32 and the shape is :math:`(n,)`.
-
-    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, atom_numbers, dt):
-        """Initialize MDIterationLeapFrog"""
-        validator.check_value_type('atom_numbers', atom_numbers, int, self.name)
-        validator.check_value_type('dt', dt, float, self.name)
-        self.atom_numbers = atom_numbers
-        self.dt = dt
-        self.add_prim_attr('atom_numbers', self.atom_numbers)
-        self.add_prim_attr('dt', self.dt)
-        self.init_prim_io_names(
-            inputs=['sqrt_mass_inverse', 'vel', 'crd', 'frc', 'acc', 'inverse_mass'],
-            outputs=['res'])
-        self.add_prim_attr('side_effect_mem', True)
-
-    def infer_shape(self, vel, crd, frc, acc, inverse_mass):
-        n = self.atom_numbers
-        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(len(inverse_mass), 1, Rel.EQ, "inverse_mass_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(inverse_mass[0], n, Rel.EQ, "inverse_mass_shape", self.name)
-        return [1,]
-
-    def infer_dtype(self, vel, crd, frc, acc, inverse_mass):
-        validator.check_tensor_dtype_valid('inverse_mass', inverse_mass, [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)
-        return mstype.float32
-
-
-class MDIterationLeapFrogWithMaxVel(PrimitiveWithInfer):
-    """
-    Leap frog algorithm to solve the Hamiltonian equations of motion with a maximum velocity limit.
-
-    .. warning::
-        This is an experimental prototype that is subject to change and/or deletion.
-
-    Args:
-        atom_numbers (int32): the number of atoms n.
-        dt (float32): the simulation time step.
-        max_velocity (float32): the maximum velocity limit.
-
-    Inputs:
-        - **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)`.
-        - **inverse_mass** (Tensor) - The inverse value of mass of each atom.
-          The data type is float32 and the shape is :math:`(n,)`.
-
-    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, atom_numbers, dt, max_velocity):
-        """Initialize MDIterationLeapFrogWithMaxVel"""
-        validator.check_value_type('atom_numbers', atom_numbers, int, self.name)
-        validator.check_value_type('dt', dt, float, self.name)
-        validator.check_value_type('max_velocity', max_velocity, float, self.name)
-        self.atom_numbers = atom_numbers
-        self.dt = dt
-        self.max_velocity = max_velocity
-
-        self.add_prim_attr('atom_numbers', self.atom_numbers)
-        self.add_prim_attr('dt', self.dt)
-        self.add_prim_attr('max_velocity', self.max_velocity)
-        self.init_prim_io_names(
-            inputs=['vel', 'crd', 'frc', 'acc', 'inverse_mass'],
-            outputs=['res'])
-        self.add_prim_attr('side_effect_mem', True)
-
-    def infer_shape(self, vel, crd, frc, acc, inverse_mass):
-        n = self.atom_numbers
-        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(len(inverse_mass), 1, Rel.EQ, "inverse_mass_dim", self.name)
-        validator.check_int(inverse_mass[0], n, Rel.EQ, "inverse_mass_shape", 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)
-        return [1,]
-
-    def infer_dtype(self, vel, crd, frc, acc, inverse_mass):
-        validator.check_tensor_dtype_valid('inverse_mass', inverse_mass, [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)
-        return mstype.float32
-
-
-class MDIterationGradientDescent(PrimitiveWithInfer):
-    """
-    Update the coordinate of each atom in the direction of potential for energy minimization.
-
-    .. warning::
-        This is an experimental prototype that is subject to change and/or deletion.
-
-    Args:
-        atom_numbers (int32): the number of atoms n.
-        learning_rate (float32): the update step length.
-
-    Inputs:
-        - **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)`.
-
-    Output:
-        - **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, atom_numbers, learning_rate):
-        """Initialize MDIterationGradientDescent"""
-        validator.check_value_type('atom_numbers', atom_numbers, int, self.name)
-        validator.check_value_type('learning_rate', learning_rate, float, self.name)
-        self.atom_numbers = atom_numbers
-        self.learning_rate = learning_rate
-        self.add_prim_attr('atom_numbers', self.atom_numbers)
-        self.add_prim_attr('learning_rate', self.learning_rate)
-        self.init_prim_io_names(
-            inputs=['crd', 'frc'],
-            outputs=['res'])
-        self.add_prim_attr('side_effect_mem', True)
-
-    def infer_shape(self, crd, frc):
-        n = self.atom_numbers
-        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(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)
-        return [1,]
-
-    def infer_dtype(self, crd, frc):
-        validator.check_tensor_dtype_valid('crd', crd, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('frc', frc, [mstype.float32], self.name)
-        return mstype.float32
-
-
-class BondForceWithAtomEnergyAndVirial(PrimitiveWithInfer):
-    """
-    Calculate bond force, harmonic potential energy and atom virial 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>`_.
-
-    .. warning::
-        This is an experimental prototype that is subject to change and/or deletion.
-
-    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),
-          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 of each atom.
-          The data type is float32 and the shape is :math:`(n, 3)`.
-        - **atom_e** (Tensor) - The energy of each atom.
-          The data type is float32 and the shape is :math:`(n,)`.
-        - **atom_virial** (Tensor) - The virial of 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 BondForceWithAtomEnergyAndVirial"""
-        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', 'atom_virial'])
-
-    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, "atom_a_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,], [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 mstype.float32, mstype.float32, mstype.float32
-
-
-class LJForceWithVirialEnergy(PrimitiveWithInfer):
-    """
-    Calculate the Lennard-Jones force, virial and atom energy 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>`_.
-
-    .. warning::
-        This is an experimental prototype that is subject to change and/or deletion.
-
-    Args:
-        atom_numbers (int32): the number of atoms, n.
-        cutoff (float32): the square value of cutoff.
-        pme_beta (float32): PME beta parameter, same as operator PMEReciprocalForce().
-        max_neighbor_numbers (int32): the max neighbor numbers, default 800.
-
-    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.
-          The number of atom pair is q. 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.
-          The number of atom pair is q. 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)`.
-        - **virial** (Tensor), The virial felt by each atom.
-          The data type is float32 and the shape is :math:`(n,)`.
-        - **atom_energy** (Tensor), The atom energy 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, max_neighbor_numbers=800):
-        """Initialize LJForceWithVirialEnergy"""
-        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)
-        validator.check_value_type('max_neighbor_numbers', max_neighbor_numbers, int, self.name)
-        self.atom_numbers = atom_numbers
-        self.cutoff = cutoff
-        self.pme_beta = pme_beta
-        self.max_neighbor_numbers = max_neighbor_numbers
-        self.init_prim_io_names(
-            inputs=['uint_crd', 'LJtype', 'charge', 'scaler', 'nl_numbers', 'nl_serial', 'd_LJ_A', 'd_LJ_B'],
-            outputs=['frc', 'virial', 'atom_energy'])
-        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)
-        self.add_prim_attr('max_neighbor_numbers', self.max_neighbor_numbers)
-
-    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]
-        m = self.max_neighbor_numbers
-        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_a), 1, Rel.EQ, "d_LJ_A_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], m, Rel.EQ, "nl_serial_shape[1]", 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 [n, 3], [n,], [n,]
-
-    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 mstype.float32, mstype.float32, mstype.float32
-
-
-class LJForceWithPMEDirectForceUpdate(PrimitiveWithInfer):
-    """
-    Calculate the Lennard-Jones force and PME direct force together for pressure.
-
-    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>`_.
-
-    .. warning::
-        This is an experimental prototype that is subject to change and/or deletion.
-
-    Args:
-        atom_numbers (int32): the number of atoms, n.
-        cutoff (float32): the square value of cutoff.
-        pme_beta (float32): PME beta parameter, same as operator PMEReciprocalForce().
-        need_update (int32): if need_update = 1, calculate the pressure, default 0.
-
-    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.
-          The number of atom pair is q. 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.
-          The number of atom pair is q. The data type is float32 and the shape is :math:`(q,)`.
-        - **beta** (Tensor) - PME beta parameter. The data type is float32 and the shape is :math:`(1,)`.
-
-    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, need_update=0):
-        """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)
-        validator.check_value_type('need_update', need_update, int, self.name)
-        self.atom_numbers = atom_numbers
-        self.cutoff = cutoff
-        self.pme_beta = pme_beta
-        self.need_update = need_update
-        self.init_prim_io_names(
-            inputs=['uint_crd', 'LJtype', 'charge', 'scaler', 'nl_numbers', 'nl_serial', 'd_LJ_A', 'd_LJ_B', 'beta'],
-            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)
-        self.add_prim_attr('need_update', self.need_update)
-
-    def infer_shape(self, uint_crd, ljtype, charge, scaler, nl_numbers, nl_serial, d_lj_a, d_lj_b, beta):
-        cls_name = self.name
-        n = self.atom_numbers
-        q = d_lj_a[0]
-        m = nl_serial[1]
-        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_a), 1, Rel.EQ, "d_LJ_A_dim", cls_name)
-        validator.check_int(len(d_lj_b), 1, Rel.EQ, "d_LJ_B_dim", cls_name)
-        validator.check_int(len(beta), 1, Rel.EQ, "beta_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], m, Rel.EQ, "nl_serial_shape[1]", 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)
-        validator.check_int(beta[0], 1, Rel.EQ, "beta_shape", cls_name)
-        return [n, 3]
-
-    def infer_dtype(self, uint_crd, ljtype, charge, scaler, nl_numbers, nl_serial, d_lj_a, d_lj_b, beta):
-        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)
-        validator.check_tensor_dtype_valid('beta', beta, [mstype.float32], self.name)
-        return mstype.float32
-
-
-class PMEReciprocalForceUpdate(PrimitiveWithInfer):
-    """
-    Calculate the reciprocal part of long-range Coulumb force using
-    PME(Particle Meshed Ewald) method for pressure. 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>`_.
-
-    .. warning::
-        This is an experimental prototype that is subject to change and/or deletion.
-
-    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
-        need_update (int32): if need_update = 1, calculate the pressure, default 0.
-
-    Inputs:
-        - **uint_crd** (Tensor) - The unsigned int coordinate 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,)`
-        - **beta** (Tensor) - The PME beta parameter to be updated in pressure calculation.
-          The data type is float32 and the shape is :math:`(1,)`
-
-    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, need_update=0):
-        """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)
-        validator.check_value_type('need_update', need_update, int, 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.need_update = need_update
-
-        self.init_prim_io_names(inputs=['uint_crd', 'charge', 'beta'],
-                                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)
-        self.add_prim_attr('need_update', self.need_update)
-
-    def infer_shape(self, uint_crd_shape, charge_shape, beta):
-        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(len(beta), 1, Rel.EQ, "beta_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)
-        validator.check_int(beta[0], 1, Rel.EQ, "beta_shape", cls_name)
-        return uint_crd_shape
-
-    def infer_dtype(self, uint_crd_type, charge_type, beta):
-        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('beta', beta, [mstype.float32], self.name)
-        return charge_type
-
-
-class PMEExcludedForceUpdate(PrimitiveWithInfer):
-    """
-    Calculate the excluded  part of long-range Coulumb force using
-    PME(Particle Meshed Ewald) method for pressure. 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>`_.
-
-    .. warning::
-        This is an experimental prototype that is subject to change and/or deletion.
-
-    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.
-        need_update (int32): if need_update = 1, calculate the pressure, default 0.
-
-    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,)`
-        - **beta** (Tensor) - The PME beta parameter to be updated in pressure calculation.
-          The data type is float32 and the shape is :math:`(1,)`
-
-    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, need_update=0):
-        """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)
-        validator.check_value_type('need_update', need_update, int, self.name)
-        self.atom_numbers = atom_numbers
-        self.excluded_numbers = excluded_numbers
-        self.beta = beta
-        self.need_update = need_update
-        self.init_prim_io_names(
-            inputs=['uint_crd', 'scaler', 'charge', 'excluded_list_start', 'excluded_list',
-                    'excluded_atom_numbers', 'beta'],
-            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)
-        self.add_prim_attr('need_update', self.need_update)
-
-    def infer_shape(self, uint_crd_shape, scaler_shape, charge_shape, excluded_list_start_shape, excluded_list_shape,
-                    excluded_atom_numbers_shape, beta):
-        cls_name = self.name
-        n = self.atom_numbers
-        e = self.excluded_numbers
-        validator.check_int(len(uint_crd_shape), 2, Rel.EQ, "uint_crd_dim", cls_name)
-        validator.check_int(len(scaler_shape), 1, Rel.EQ, "scaler_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(len(excluded_list_shape), 1, Rel.EQ, "excluded_list_dim", cls_name)
-        validator.check_int(len(beta), 1, Rel.EQ, "beta_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(scaler_shape[0], 3, Rel.EQ, "scaler_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)
-        validator.check_int(excluded_list_shape[0], e, Rel.EQ, "excluded_list_shape", cls_name)
-        validator.check_int(beta[0], 1, Rel.EQ, "beta_shape", cls_name)
-        return [n, 3]
-
-    def infer_dtype(self, uint_crd_type, scaler_type, charge_type, excluded_list_start_type, excluded_list_type,
-                    excluded_atom_numbers_type, beta):
-        validator.check_tensor_dtype_valid('scaler', scaler_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)
-        validator.check_tensor_dtype_valid('beta', beta, [mstype.float32], self.name)
-        return mstype.float32
-
-
-class LJForceWithVirialEnergyUpdate(PrimitiveWithInfer):
-    """
-    Calculate the Lennard-Jones force and PME direct force together for pressure.
-
-    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>`_.
-
-    .. warning::
-        This is an experimental prototype that is subject to change and/or deletion.
-
-    Args:
-        atom_numbers (int32): the number of atoms, n.
-        cutoff (float32): the square value of cutoff.
-        pme_beta (float32): PME beta parameter, same as operator PMEReciprocalForce().
-        max_neighbor_numbers (int32): the max neighbor numbers, default 800.
-        need_update (int32): if need_update = 1, calculate the pressure, default 0.
-
-    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.
-          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.
-          The number of atom pair is q. 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.
-          The number of atom pair is q. The data type is float32 and the shape is :math:`(q,)`.
-        - **beta** (Tensor) - The PME beta parameter to be updated in pressure calculation.
-          The data type is float32 and the shape is :math:`(1,)`
-
-    Outputs:
-        - **frc** (Tensor) - The force felt by each atom.
-          The data type is float32 and the shape is :math:`(n, 3)`.
-        - **virial** (Tensor) - The accumulated potential virial for each atom.
-          The data type is float32 and the shape is :math:`(n, )`.
-        - **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, atom_numbers, cutoff, pme_beta, max_neighbor_numbers=800, need_update=0):
-        """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)
-        validator.check_value_type('max_neighbor_numbers', max_neighbor_numbers, int, self.name)
-        validator.check_value_type('need_update', need_update, int, self.name)
-        self.atom_numbers = atom_numbers
-        self.cutoff = cutoff
-        self.pme_beta = pme_beta
-        self.max_neighbor_numbers = max_neighbor_numbers
-        self.need_update = need_update
-        self.init_prim_io_names(
-            inputs=['uint_crd', 'LJtype', 'charge', 'scaler', 'nl_numbers', 'nl_serial', 'd_LJ_A', 'd_LJ_B', 'beta'],
-            outputs=['frc', 'virial', 'atom_energy'])
-        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)
-        self.add_prim_attr('max_neighbor_numbers', self.max_neighbor_numbers)
-        self.add_prim_attr('need_update', self.need_update)
-
-    def infer_shape(self, uint_crd, ljtype, charge, scaler, nl_numbers, nl_serial, d_lj_a, d_lj_b, beta):
-        cls_name = self.name
-        n = self.atom_numbers
-        q = d_lj_a[0]
-        m = self.max_neighbor_numbers
-        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_a), 1, Rel.EQ, "d_LJ_A_dim", cls_name)
-        validator.check_int(len(d_lj_b), 1, Rel.EQ, "d_LJ_B_dim", cls_name)
-        validator.check_int(len(beta), 1, Rel.EQ, "beta_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], m, Rel.EQ, "nl_serial_shape[1]", 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)
-        validator.check_int(beta[0], 1, Rel.EQ, "beta_shape[0]", cls_name)
-        return [n, 3], [n,], [n,]
-
-    def infer_dtype(self, uint_crd, ljtype, charge, scaler, nl_numbers, nl_serial, d_lj_a, d_lj_b, beta):
-        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)
-        validator.check_tensor_dtype_valid('beta', beta, [mstype.float32], self.name)
-        return mstype.float32, mstype.float32, mstype.float32
-
-
-class Dihedral14ForceWithAtomEnergyVirial(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`.
-
-    .. warning::
-        This is an experimental prototype that is subject to change and/or deletion.
-
-    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)`.
-        - **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,)`.
-        - **boxlength** (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.
-        - **cf_scale_factor** (Tensor) - The scale factor for the Coulomb force.
-          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.
-          The number of atom pair is q. 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.
-          The number of atom pair is q. 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)`.
-        - **atom_energy** (Tensor) - The accumulated potential energy for each atom.
-          The data type is float32 and the shape is :math:`(n, )`.
-        - **atom_virial** (Tensor) - The accumulated potential virial 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', 'a_14', 'b_14', 'lj_scale_factor',
-                    'cf_scale_factor', 'LJ_type_A', 'LJ_type_B'],
-            outputs=['frc', 'atom_energy', 'atom_virial'])
-        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_a_shape), 1, Rel.EQ, "LJ_type_A_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_a_shape[0], q, Rel.EQ, "LJ_type_A_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 [n, 3], [n,], [n,]
-
-    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 mstype.float32, mstype.float32, mstype.float32
-
-
-class PMEEnergyUpdate(PrimitiveWithInfer):
-    """
-    Calculate the Coulumb energy of the system using PME method for pressure.
-
-    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>`_.
-
-    .. warning::
-        This is an experimental prototype that is subject to change and/or deletion.
-
-    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.
-        max_neighbor_numbers (int32): the max neighbor numbers, m, default 800.
-        need_update (int32): if need_update = 1, calculate the pressure, default 0.
-
-    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, m)`
-        - **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,)`
-        - **factor** (Tensor) - The factor parameter to be updated in pressure calculation.
-          The data type is float32 and the shape is :math:`(1,)`
-        - **beta** (Tensor) - The PME beta parameter to be updated in pressure calculation.
-          The data type is float32 and the shape is :math:`(1,)`
-
-    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, max_neighbor_numbers=800, need_update=0):
-        """Initialize PMEEnergyUpdate"""
-        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)
-        validator.check_value_type('max_neighbor_numbers', max_neighbor_numbers, int, self.name)
-        validator.check_value_type('need_update', need_update, int, 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.max_neighbor_numbers = max_neighbor_numbers
-        self.need_update = need_update
-        self.init_prim_io_names(
-            inputs=['uint_crd', 'charge', 'nl_numbers', 'nl_serial', 'scaler', 'excluded_list_start',
-                    'excluded_list', 'excluded_atom_numbers', 'factor', 'beta'],
-            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)
-        self.add_prim_attr('max_neighbor_numbers', self.max_neighbor_numbers)
-        self.add_prim_attr('need_update', self.need_update)
-
-    def infer_shape(self, uint_crd, charge, nl_numbers, nl_serial, scaler, excluded_list_start,
-                    excluded_list, excluded_atom_numbers, factor, beta):
-        cls_name = self.name
-        n = self.atom_numbers
-        m = self.max_neighbor_numbers
-        e = self.excluded_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.EQ, "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.EQ, "excluded_list_dim", cls_name)
-        validator.check_int(len(scaler), 1, Rel.EQ, "scaler_dim", cls_name)
-        validator.check_int(len(factor), 1, Rel.EQ, "factor_dim", cls_name)
-        validator.check_int(len(beta), 1, Rel.EQ, "beta_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(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.EQ, "nl_serial_shape[0]", cls_name)
-        validator.check_int(nl_serial[1], m, Rel.EQ, "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], e, Rel.EQ, "excluded_list_shape", cls_name)
-        validator.check_int(factor[0], 1, Rel.EQ, "factor_shape", cls_name)
-        validator.check_int(beta[0], 1, Rel.EQ, "beta_shape", cls_name)
-        validator.check_int(scaler[0], 3, Rel.EQ, "scaler_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, factor, beta):
-        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)
-        validator.check_tensor_dtype_valid('factor', factor, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('beta', beta, [mstype.float32], self.name)
-        return charge, charge, charge, charge
-
-
-class ConstrainForceCycle(PrimitiveWithInfer):
-    """
-    Calculate the constraint force in each iteration.
-
-    .. warning::
-        This is an experimental prototype that is subject to change and/or deletion.
-
-    Args:
-        atom_numbers (int32): the number of atoms n.
-        constrain_pair_numbers (int32): the number of constrain pairs m.
-
-    Inputs:
-        - **uint_crd** (Tensor) - The unsigned int coordinate value of each atom.
-          The data type is uint32 and the shape is :math:`(n, 3)`.
-        - **scaler** (Tensor) - The 3-D scale factor (x, y, z),
-          The data type is float32 and the shape is :math:`(3,)`.
-        - **pair_dr** (Tensor) - The displacement vector of each constrained atom pair.
-          The data type is float32 and the shape is :math:`(m, 3)`.
-        - **atom_i_serials** (Tensor) - The first atom index of each constrained atom pair.
-          The data type is int32 and the shape is :math:`(m,)`.
-        - **atom_j_serials** (Tensor) - The second atom index of each constrained atom pair.
-          The data type is int32 and the shape is :math:`(m,)`.
-        - **constant_rs** (Tensor) - The constrained distance of each constrained atom pair.
-          The data type is float32 and the shape is :math:`(m,)`.
-        - **constrain_ks** (Tensor) - The coefficient of each constrained atom pair.
-          The data type is float32 and the shape is :math:`(m,)`.
-
-    Outputs:
-        - **test_frc** (Tensor) - The constraint force.
-          The data type is float32 and the shape is :math:`(n, 3)`.
-
-    Supported Platforms:
-        ``GPU``
-    """
-
-    @prim_attr_register
-    def __init__(self, atom_numbers, constrain_pair_numbers):
-        validator.check_value_type('atom_numbers', atom_numbers, int, self.name)
-        validator.check_value_type('constrain_pair_numbers', constrain_pair_numbers, int, self.name)
-        self.atom_numbers = atom_numbers
-        self.constrain_pair_numbers = constrain_pair_numbers
-        self.add_prim_attr('atom_numbers', self.atom_numbers)
-        self.add_prim_attr('constrain_pair_numbers', self.constrain_pair_numbers)
-        self.init_prim_io_names(
-            inputs=['uint_crd', 'scaler', 'pair_dr', 'atom_i_serials', 'atom_j_serials',
-                    'constant_rs', 'constrain_ks'],
-            outputs=['test_frc'])
-
-    def infer_shape(self, uint_crd_shape, scaler_shape, pair_dr_shape, atom_i_serials_shape,
-                    atom_j_serials_shape, constant_rs_shape, constrain_ks_shape):
-        cls_name = self.name
-        n = self.atom_numbers
-        m = self.constrain_pair_numbers
-        validator.check_int(len(uint_crd_shape), 2, Rel.EQ, "uint_crd_dim", cls_name)
-        validator.check_int(len(scaler_shape), 1, Rel.EQ, "scaler_dim", cls_name)
-        validator.check_int(len(pair_dr_shape), 2, Rel.EQ, "pair_dr_dim", cls_name)
-        validator.check_int(len(atom_i_serials_shape), 1, Rel.EQ, "atom_i_serials_dim", cls_name)
-        validator.check_int(len(atom_j_serials_shape), 1, Rel.EQ, "atom_j_serials_dim", cls_name)
-        validator.check_int(len(constant_rs_shape), 1, Rel.EQ, "constant_rs_dim", cls_name)
-        validator.check_int(len(constrain_ks_shape), 1, Rel.EQ, "constrain_ks_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(scaler_shape[0], 3, Rel.EQ, "scaler_shape", cls_name)
-        validator.check_int(pair_dr_shape[0], m, Rel.EQ, "pair_dr_shape[0]", cls_name)
-        validator.check_int(pair_dr_shape[1], 3, Rel.EQ, "pair_dr_shape[1]", cls_name)
-        validator.check_int(atom_i_serials_shape[0], m, Rel.EQ, "atom_i_serials_shape[0]", cls_name)
-        validator.check_int(atom_j_serials_shape[0], m, Rel.EQ, "atom_j_serials_shape[0]", cls_name)
-        validator.check_int(constant_rs_shape[0], m, Rel.EQ, "constant_rs_shape[0]", cls_name)
-        validator.check_int(constrain_ks_shape[0], m, Rel.EQ, "constrain_ks_shape[0]", cls_name)
-        return [n, 3]
-
-    def infer_dtype(self, uint_crd_dtype, scaler_dtype, pair_dr_dtype, atom_i_serials_dtype,
-                    atom_j_serials_dtype, constant_rs_dtype, constrain_ks_dtype):
-        validator.check_tensor_dtype_valid('uint_crd', uint_crd_dtype, [mstype.uint32], self.name)
-        validator.check_tensor_dtype_valid('scaler', scaler_dtype, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('pair_dr', pair_dr_dtype, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('atom_i_serials', atom_i_serials_dtype, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('atom_j_serials', atom_j_serials_dtype, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('constant_rs', constant_rs_dtype, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('constrain_ks', constrain_ks_dtype, [mstype.float32], self.name)
-        return mstype.float32
-
-
-class ConstrainForceVirial(PrimitiveWithInfer):
-    """
-    Calculate the constraint force and virial in a step with iteration numbers.
-
-    .. warning::
-        This is an experimental prototype that is subject to change and/or deletion.
-
-    Args:
-        atom_numbers (int32): the number of atoms n.
-        constrain_pair_numbers (int32): the number of constrain pairs m.
-        iteration_numbers (int32): the number of iteration numbers p.
-        half_exp_gamma_plus_half (float32): half exp_gamma plus half q.
-
-    Inputs:
-        - **crd** (Tensor) - The coordinate of each atom.
-          The data type is float32 and the shape is :math:`(n, 3)`.
-        - **quarter_cof** (Tensor) - The 3-D scale factor.
-          The data type is float32 and the shape is :math:`(3,)`.
-        - **mass_inverse** (Tensor) - The inverse value of mass of each atom.
-          The data type is float32 and the shape is :math:`(n,)`.
-        - **scaler** (Tensor) - The 3-D scale factor (x, y, z),
-          The data type is float32 and the shape is :math:`(3,)`.
-        - **pair_dr** (Tensor) - The displacement vector of each constrained atom pair.
-          The data type is float32 and the shape is :math:`(m, 3)`.
-        - **atom_i_serials** (Tensor) - The first atom index of each constrained atom pair.
-          The data type is int32 and the shape is :math:`(m,)`.
-        - **atom_j_serials** (Tensor) - The second atom index of each constrained atom pair.
-          The data type is int32 and the shape is :math:`(m,)`.
-        - **constant_rs** (Tensor) - The constrained distance of each constrained atom pair.
-          The data type is float32 and the shape is :math:`(m,)`.
-        - **constrain_ks** (Tensor) - The coefficient of each constrained atom pair.
-          The data type is float32 and the shape is :math:`(m,)`.
-
-    Outputs:
-        - **uint_crd** (Tensor) - The unsigned int coordinate value of each atom.
-          The data type is uint32 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)`.
-        - **virial** (Tensor) - The constraint virial on each atom.
-          The data type is float32 and the shape is :math:`(m,)`.
-
-    Supported Platforms:
-        ``GPU``
-    """
-
-    @prim_attr_register
-    def __init__(self, atom_numbers, constrain_pair_numbers, iteration_numbers, half_exp_gamma_plus_half):
-        validator.check_value_type('atom_numbers', atom_numbers, int, self.name)
-        validator.check_value_type('constrain_pair_numbers', constrain_pair_numbers, int, self.name)
-        validator.check_value_type('iteration_numbers', iteration_numbers, int, self.name)
-        validator.check_value_type('half_exp_gamma_plus_half', half_exp_gamma_plus_half, float, self.name)
-        self.atom_numbers = atom_numbers
-        self.constrain_pair_numbers = constrain_pair_numbers
-        self.iteration_numbers = iteration_numbers
-        self.half_exp_gamma_plus_half = half_exp_gamma_plus_half
-        self.add_prim_attr('atom_numbers', self.atom_numbers)
-        self.add_prim_attr('constrain_pair_numbers', self.constrain_pair_numbers)
-        self.add_prim_attr('iteration_numbers', self.iteration_numbers)
-        self.add_prim_attr('half_exp_gamma_plus_half', self.half_exp_gamma_plus_half)
-
-        self.init_prim_io_names(
-            inputs=['crd', 'quarter_cof', 'mass_inverse',
-                    'scaler', 'pair_dr', 'atom_i_serials', 'atom_j_serials',
-                    'constant_rs', 'constrain_ks'],
-            outputs=['uint_crd', 'frc', 'virial'])
-
-    def infer_shape(self, crd, quarter_cof, mass_inverse, scaler_shape, pair_dr_shape, atom_i_serials_shape,
-                    atom_j_serials_shape, constant_rs_shape, constrain_ks_shape):
-        cls_name = self.name
-        n = self.atom_numbers
-        m = self.constrain_pair_numbers
-        validator.check_int(len(crd), 2, Rel.EQ, "crd_dim", cls_name)
-        validator.check_int(len(quarter_cof), 1, Rel.EQ, "quarter_cof_dim", cls_name)
-        validator.check_int(len(mass_inverse), 1, Rel.EQ, "mass_inverse_dim", cls_name)
-        validator.check_int(len(scaler_shape), 1, Rel.EQ, "scaler_dim", cls_name)
-        validator.check_int(len(pair_dr_shape), 2, Rel.EQ, "pair_dr_dim", cls_name)
-        validator.check_int(len(atom_i_serials_shape), 1, Rel.EQ, "atom_i_serials_dim", cls_name)
-        validator.check_int(len(atom_j_serials_shape), 1, Rel.EQ, "atom_j_serials_dim", cls_name)
-        validator.check_int(len(constant_rs_shape), 1, Rel.EQ, "constant_rs_dim", cls_name)
-        validator.check_int(len(constrain_ks_shape), 1, Rel.EQ, "constrain_ks_dim", cls_name)
-        validator.check_int(crd[0], n, Rel.EQ, "crd_shape[0]", cls_name)
-        validator.check_int(crd[1], 3, Rel.EQ, "crd_shape[1]", cls_name)
-        validator.check_int(quarter_cof[0], 3, Rel.EQ, "quarter_cof_shape", cls_name)
-        validator.check_int(mass_inverse[0], n, Rel.EQ, "mass_inverse_shape", cls_name)
-        validator.check_int(scaler_shape[0], 3, Rel.EQ, "scaler_shape", cls_name)
-        validator.check_int(pair_dr_shape[0], m, Rel.EQ, "pair_dr_shape[0]", cls_name)
-        validator.check_int(pair_dr_shape[1], 3, Rel.EQ, "pair_dr_shape[1]", cls_name)
-        validator.check_int(atom_i_serials_shape[0], m, Rel.EQ, "atom_i_serials_shape[0]", cls_name)
-        validator.check_int(atom_j_serials_shape[0], m, Rel.EQ, "atom_j_serials_shape[0]", cls_name)
-        validator.check_int(constant_rs_shape[0], m, Rel.EQ, "constant_rs_shape[0]", cls_name)
-        validator.check_int(constrain_ks_shape[0], m, Rel.EQ, "constrain_ks_shape[0]", cls_name)
-        return [n, 3], [n, 3], [m,]
-
-    def infer_dtype(self, crd, quarter_cof, mass_inverse, scaler_dtype, pair_dr_dtype, atom_i_serials_dtype,
-                    atom_j_serials_dtype, constant_rs_dtype, constrain_ks_dtype):
-        validator.check_tensor_dtype_valid('crd', crd, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('quarter_cof', quarter_cof, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('mass_inverse', mass_inverse, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('scaler', scaler_dtype, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('pair_dr', pair_dr_dtype, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('atom_i_serials', atom_i_serials_dtype, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('atom_j_serials', atom_j_serials_dtype, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('constant_rs', constant_rs_dtype, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('constrain_ks', constrain_ks_dtype, [mstype.float32], self.name)
-        return mstype.uint32, mstype.float32, mstype.float32
-
-
-class ConstrainForce(PrimitiveWithInfer):
-    """
-    Calculate the constraint force in a step with iteration numbers.
-
-    .. warning::
-        This is an experimental prototype that is subject to change and/or deletion.
-
-    Args:
-        atom_numbers (int32): the number of atoms n.
-        constrain_pair_numbers (int32): the number of constrain pairs m.
-        iteration_numbers (int32): the number of iteration numbers p.
-        half_exp_gamma_plus_half (float32): half exp_gamma plus half q.
-
-    Inputs:
-        - **crd** (Tensor) - The coordinate of each atom.
-          The data type is float32 and the shape is :math:`(n, 3)`.
-        - **quarter_cof** (Tensor) - The 3-D scale factor.
-          The data type is float32 and the shape is :math:`(3,)`.
-        - **mass_inverse** (Tensor) - The inverse value of mass of each atom.
-          The data type is float32 and the shape is :math:`(n,)`.
-        - **scaler** (Tensor) - The 3-D scale factor (x, y, z),
-          The data type is float32 and the shape is :math:`(3,)`.
-        - **pair_dr** (Tensor) - The displacement vector of each constrained atom pair.
-          The data type is float32 and the shape is :math:`(m, 3)`.
-        - **atom_i_serials** (Tensor) - The first atom index of each constrained atom pair.
-          The data type is int32 and the shape is :math:`(m,)`.
-        - **atom_j_serials** (Tensor) - The second atom index of each constrained atom pair.
-          The data type is int32 and the shape is :math:`(m,)`.
-        - **constant_rs** (Tensor) - The constrained distance of each constrained atom pair.
-          The data type is float32 and the shape is :math:`(m,)`.
-        - **constrain_ks** (Tensor) - The coefficient of each constrained atom pair.
-          The data type is float32 and the shape is :math:`(m,)`.
-
-    Outputs:
-        - **uint_crd** (Tensor) - The unsigned int coordinate value of each atom.
-          The data type is uint32 and the shape is :math:`(n, 3)`.
-        - **frc** (Tensor) - The constraint force on each atom.
-          The data type is float32 and the shape is :math:`(n, 3)`.
-        - **virial** (Tensor) - The constraint virial on each atom and it is zero.
-          The data type is float32 and the shape is :math:`(m,)`.
-
-    Supported Platforms:
-        ``GPU``
-    """
-
-    @prim_attr_register
-    def __init__(self, atom_numbers, constrain_pair_numbers, iteration_numbers, half_exp_gamma_plus_half):
-        validator.check_value_type('atom_numbers', atom_numbers, int, self.name)
-        validator.check_value_type('constrain_pair_numbers', constrain_pair_numbers, int, self.name)
-        validator.check_value_type('iteration_numbers', iteration_numbers, int, self.name)
-        validator.check_value_type('half_exp_gamma_plus_half', half_exp_gamma_plus_half, float, self.name)
-        self.atom_numbers = atom_numbers
-        self.constrain_pair_numbers = constrain_pair_numbers
-        self.iteration_numbers = iteration_numbers
-        self.half_exp_gamma_plus_half = half_exp_gamma_plus_half
-        self.add_prim_attr('atom_numbers', self.atom_numbers)
-        self.add_prim_attr('constrain_pair_numbers', self.constrain_pair_numbers)
-        self.add_prim_attr('iteration_numbers', self.iteration_numbers)
-        self.add_prim_attr('half_exp_gamma_plus_half', self.half_exp_gamma_plus_half)
-
-        self.init_prim_io_names(
-            inputs=['crd', 'quarter_cof', 'mass_inverse',
-                    'scaler', 'pair_dr', 'atom_i_serials', 'atom_j_serials', 'constant_rs', 'constrain_ks'],
-            outputs=['uint_crd', 'frc', 'virial'])
-
-    def infer_shape(self, crd, quarter_cof, mass_inverse, scaler_shape, pair_dr_shape, atom_i_serials_shape,
-                    atom_j_serials_shape, constant_rs_shape, constrain_ks_shape):
-        cls_name = self.name
-        n = self.atom_numbers
-        m = self.constrain_pair_numbers
-        validator.check_int(len(crd), 2, Rel.EQ, "crd_dim", cls_name)
-        validator.check_int(len(quarter_cof), 1, Rel.EQ, "quarter_cof_dim", cls_name)
-        validator.check_int(len(mass_inverse), 1, Rel.EQ, "mass_inverse_dim", cls_name)
-        validator.check_int(len(scaler_shape), 1, Rel.EQ, "scaler_dim", cls_name)
-        validator.check_int(len(pair_dr_shape), 2, Rel.EQ, "pair_dr_dim", cls_name)
-        validator.check_int(len(atom_i_serials_shape), 1, Rel.EQ, "atom_i_serials_dim", cls_name)
-        validator.check_int(len(atom_j_serials_shape), 1, Rel.EQ, "atom_j_serials_dim", cls_name)
-        validator.check_int(len(constant_rs_shape), 1, Rel.EQ, "constant_rs_dim", cls_name)
-        validator.check_int(len(constrain_ks_shape), 1, Rel.EQ, "constrain_ks_dim", cls_name)
-        validator.check_int(crd[0], n, Rel.EQ, "crd_shape[0]", cls_name)
-        validator.check_int(crd[1], 3, Rel.EQ, "crd_shape[1]", cls_name)
-        validator.check_int(quarter_cof[0], 3, Rel.EQ, "quarter_cof_shape", cls_name)
-        validator.check_int(mass_inverse[0], n, Rel.EQ, "mass_inverse_shape", cls_name)
-        validator.check_int(scaler_shape[0], 3, Rel.EQ, "scaler_shape", cls_name)
-        validator.check_int(pair_dr_shape[0], m, Rel.EQ, "pair_dr_shape[0]", cls_name)
-        validator.check_int(pair_dr_shape[1], 3, Rel.EQ, "pair_dr_shape[1]", cls_name)
-        validator.check_int(atom_i_serials_shape[0], m, Rel.EQ, "atom_i_serials_shape[0]", cls_name)
-        validator.check_int(atom_j_serials_shape[0], m, Rel.EQ, "atom_j_serials_shape[0]", cls_name)
-        validator.check_int(constant_rs_shape[0], m, Rel.EQ, "constant_rs_shape[0]", cls_name)
-        validator.check_int(constrain_ks_shape[0], m, Rel.EQ, "constrain_ks_shape[0]", cls_name)
-        return [n, 3], [n, 3], [m,]
-
-    def infer_dtype(self, crd, quarter_cof, mass_inverse, scaler_dtype, pair_dr_dtype, atom_i_serials_dtype,
-                    atom_j_serials_dtype, constant_rs_dtype, constrain_ks_dtype):
-        validator.check_tensor_dtype_valid('crd', crd, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('quarter_cof', quarter_cof, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('mass_inverse', mass_inverse, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('scaler', scaler_dtype, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('pair_dr', pair_dr_dtype, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('atom_i_serials', atom_i_serials_dtype, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('atom_j_serials', atom_j_serials_dtype, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('constant_rs', constant_rs_dtype, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('constrain_ks', constrain_ks_dtype, [mstype.float32], self.name)
-        return mstype.uint32, mstype.float32, mstype.float32
-
-
-class Constrain(PrimitiveWithInfer):
-    """
-    Calculate the constraint force and virial depends on pressure calculation.
-
-    .. warning::
-        This is an experimental prototype that is subject to change and/or deletion.
-
-    Args:
-        atom_numbers (int32): the number of atoms n.
-        constrain_pair_numbers (int32): the number of constrain pairs m.
-        iteration_numbers (int32): the number of iteration numbers p.
-        half_exp_gamma_plus_half (float32): half exp_gamma plus half q.
-        update_interval (int32): the number of update interval, default 10.
-
-    Inputs:
-        - **crd** (Tensor) - The coordinate of each atom.
-          The data type is float32 and the shape is :math:`(n, 3)`.
-        - **quarter_cof** (Tensor) - The 3-D scale factor.
-          The data type is float32 and the shape is :math:`(3,)`.
-        - **mass_inverse** (Tensor) - The inverse value of mass of each atom.
-          The data type is float32 and the shape is :math:`(n,)`.
-        - **scaler** (Tensor) - The 3-D scale factor (x, y, z),
-          The data type is float32 and the shape is :math:`(3,)`.
-        - **pair_dr** (Tensor) - The displacement vector of each constrained atom pair.
-          The data type is float32 and the shape is :math:`(m, 3)`.
-        - **atom_i_serials** (Tensor) - The first atom index of each constrained atom pair.
-          The data type is int32 and the shape is :math:`(m,)`.
-        - **atom_j_serials** (Tensor) - The second atom index of each constrained atom pair.
-          The data type is int32 and the shape is :math:`(m,)`.
-        - **constant_rs** (Tensor) - The constrained distance of each constrained atom pair.
-          The data type is float32 and the shape is :math:`(m,)`.
-        - **constrain_ks** (Tensor) - The coefficient of each constrained atom pair.
-          The data type is float32 and the shape is :math:`(m,)`.
-        - **need_pressure** (Tensor) - If need pressure, 1 else 0.
-          The data type is int32 and the shape is :math:`(1,)` or :math:`()`.
-
-    Outputs:
-        - **uint_crd** (Tensor) - The unsigned int coordinate value of each atom.
-          The data type is uint32 and the shape is :math:`(n, 3)`.
-        - **frc** (Tensor) - The constraint force on each atom.
-          The data type is float32 and the shape is :math:`(n, 3)`.
-        - **virial** (Tensor) - The constraint virial on each atom.
-          The data type is float32 and the shape is :math:`(m,)`.
-
-    Supported Platforms:
-        ``GPU``
-    """
-
-    @prim_attr_register
-    def __init__(self, atom_numbers, constrain_pair_numbers, iteration_numbers, half_exp_gamma_plus_half,
-                 update_interval=10):
-        validator.check_value_type('atom_numbers', atom_numbers, int, self.name)
-        validator.check_value_type('constrain_pair_numbers', constrain_pair_numbers, int, self.name)
-        validator.check_value_type('iteration_numbers', iteration_numbers, int, self.name)
-        validator.check_value_type('half_exp_gamma_plus_half', half_exp_gamma_plus_half, float, self.name)
-        validator.check_value_type('update_interval', update_interval, int, self.name)
-        self.atom_numbers = atom_numbers
-        self.constrain_pair_numbers = constrain_pair_numbers
-        self.iteration_numbers = iteration_numbers
-        self.half_exp_gamma_plus_half = half_exp_gamma_plus_half
-        self.update_interval = update_interval
-        self.add_prim_attr('atom_numbers', self.atom_numbers)
-        self.add_prim_attr('constrain_pair_numbers', self.constrain_pair_numbers)
-        self.add_prim_attr('iteration_numbers', self.iteration_numbers)
-        self.add_prim_attr('half_exp_gamma_plus_half', self.half_exp_gamma_plus_half)
-        self.add_prim_attr('update_interval', self.update_interval)
-
-        self.init_prim_io_names(
-            inputs=['crd', 'quarter_cof', 'mass_inverse',
-                    'scaler', 'pair_dr', 'atom_i_serials', 'atom_j_serials',
-                    'constant_rs', 'constrain_ks', 'need_pressure'],
-            outputs=['uint_crd', 'frc', 'virial'])
-
-    def infer_shape(self, crd, quarter_cof, mass_inverse, scaler_shape, pair_dr_shape, atom_i_serials_shape,
-                    atom_j_serials_shape, constant_rs_shape, constrain_ks_shape, need_pressure):
-        cls_name = self.name
-        n = self.atom_numbers
-        m = self.constrain_pair_numbers
-        validator.check_int(len(crd), 2, Rel.EQ, "crd_dim", cls_name)
-        validator.check_int(len(quarter_cof), 1, Rel.EQ, "quarter_cof_dim", cls_name)
-        validator.check_int(len(mass_inverse), 1, Rel.EQ, "mass_inverse_dim", cls_name)
-        validator.check_int(len(scaler_shape), 1, Rel.EQ, "scaler_dim", cls_name)
-        validator.check_int(len(pair_dr_shape), 2, Rel.EQ, "pair_dr_dim", cls_name)
-        validator.check_int(len(atom_i_serials_shape), 1, Rel.EQ, "atom_i_serials_dim", cls_name)
-        validator.check_int(len(atom_j_serials_shape), 1, Rel.EQ, "atom_j_serials_dim", cls_name)
-        validator.check_int(len(constant_rs_shape), 1, Rel.EQ, "constant_rs_dim", cls_name)
-        validator.check_int(len(constrain_ks_shape), 1, Rel.EQ, "constrain_ks_dim", cls_name)
-        validator.check_int(len(need_pressure), 1, Rel.LE, "need_pressure_dim", cls_name)
-        validator.check_int(crd[0], n, Rel.EQ, "crd_shape[0]", cls_name)
-        validator.check_int(crd[1], 3, Rel.EQ, "crd_shape[1]", cls_name)
-        validator.check_int(quarter_cof[0], 3, Rel.EQ, "quarter_cof_shape", cls_name)
-        validator.check_int(mass_inverse[0], n, Rel.EQ, "mass_inverse_shape", cls_name)
-        validator.check_int(scaler_shape[0], 3, Rel.EQ, "scaler_shape", cls_name)
-        validator.check_int(pair_dr_shape[0], m, Rel.EQ, "pair_dr_shape[0]", cls_name)
-        validator.check_int(pair_dr_shape[1], 3, Rel.EQ, "pair_dr_shape[1]", cls_name)
-        validator.check_int(atom_i_serials_shape[0], m, Rel.EQ, "atom_i_serials_shape[0]", cls_name)
-        validator.check_int(atom_j_serials_shape[0], m, Rel.EQ, "atom_j_serials_shape[0]", cls_name)
-        validator.check_int(constant_rs_shape[0], m, Rel.EQ, "constant_rs_shape[0]", cls_name)
-        validator.check_int(constrain_ks_shape[0], m, Rel.EQ, "constrain_ks_shape[0]", cls_name)
-        if need_pressure:
-            validator.check_int(need_pressure[0], 1, Rel.EQ, "need_pressure_shape", self.name)
-        return [n, 3], [n, 3], [m,]
-
-    def infer_dtype(self, crd, quarter_cof, mass_inverse, scaler_dtype, pair_dr_dtype, atom_i_serials_dtype,
-                    atom_j_serials_dtype, constant_rs_dtype, constrain_ks_dtype, need_pressure):
-        validator.check_tensor_dtype_valid('crd', crd, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('quarter_cof', quarter_cof, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('mass_inverse', mass_inverse, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('scaler', scaler_dtype, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('pair_dr', pair_dr_dtype, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('atom_i_serials', atom_i_serials_dtype, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('atom_j_serials', atom_j_serials_dtype, [mstype.int32], self.name)
-        validator.check_tensor_dtype_valid('constant_rs', constant_rs_dtype, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('constrain_ks', constrain_ks_dtype, [mstype.float32], self.name)
-        validator.check_tensor_dtype_valid('need_pressure', need_pressure, [mstype.int32], self.name)
-        return mstype.uint32, mstype.float32, mstype.float32