From e1c277cdc5129ac1442b40577e42aad8c061ee99 Mon Sep 17 00:00:00 2001 From: zhangxinfeng3 <zhangxinfeng3@huawei.com> Date: Wed, 3 Mar 2021 15:03:01 +0800 Subject: [PATCH] update sponge --- mindspore/ops/operations/sponge_ops.py | 324 ++++++++++++++++++++++--- 1 file changed, 294 insertions(+), 30 deletions(-) diff --git a/mindspore/ops/operations/sponge_ops.py b/mindspore/ops/operations/sponge_ops.py index 1b2830cf1bc..3737fce0436 100644 --- a/mindspore/ops/operations/sponge_ops.py +++ b/mindspore/ops/operations/sponge_ops.py @@ -18,6 +18,7 @@ from ..primitive import PrimitiveWithInfer, prim_attr_register from ..._checkparam import Validator as validator from ...common import dtype as mstype +from ..._checkparam import Rel class BondForce(PrimitiveWithInfer): @@ -50,12 +51,30 @@ class BondForce(PrimitiveWithInfer): ``GPU`` Examples: """ + @prim_attr_register def __init__(self, bond_numbers): self.bond_numbers = bond_numbers self.init_prim_io_names(inputs=['uint_crd_f', 'scaler_f', 'atom_a', 'atom_b', 'bond_k', 'bond_r0'], outputs=['frc_f']) self.add_prim_attr('bond_numbers', self.bond_numbers) + + 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 = uint_crd_f_shape[0] + M = atom_a_shape[0] + validator.check_int( + uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape", cls_name) + validator.check_int( + scaler_f_shape[0], 3, Rel.EQ, "scaler_f_shape", cls_name) + validator.check_int( + atom_b_shape[0], M, Rel.EQ, "atom_b_shape", cls_name) + validator.check_int( + bond_k_shape[0], M, Rel.EQ, "bond_k_shape", cls_name) + validator.check_int( + bond_r0_shape[0], M, Rel.EQ, "bond_r0_shape", cls_name) + return uint_crd_f_shape + def infer_dtype(self, uint_crd_f_dtype, scaler_f_type, atom_a_type, atom_b_type, bond_k_type, bond_r0_type): validator.check_tensor_dtype_valid('uint_crd_f_dtype', uint_crd_f_dtype, [mstype.uint32], self.name) validator.check_tensor_dtype_valid('scaler_f_type', scaler_f_type, [mstype.float32], self.name) @@ -81,6 +100,11 @@ class BondEnergy(PrimitiveWithInfer): Inputs: Same as operator BondForce(). + .. math:: + + dr = (x_1-x_2, y_1-y_2, z_1-z_2) + E = k*(|dr| - r_0)^2 + Outputs: - **bond_ene** (Tensor, float32) - [M, 1], the harmonic potential energy for each bond. @@ -89,12 +113,31 @@ class BondEnergy(PrimitiveWithInfer): ``GPU`` Examples: """ + @prim_attr_register def __init__(self, bond_numbers): self.bond_numbers = bond_numbers self.init_prim_io_names(inputs=['uint_crd_f', 'scaler_f', 'atom_a', 'atom_b', 'bond_k', 'bond_r0'], outputs=['bond_ene']) self.add_prim_attr('bond_numbers', self.bond_numbers) + + 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 = uint_crd_f_shape[0] + M = atom_a_shape[0] + validator.check_int( + uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape", cls_name) + validator.check_int( + scaler_f_shape[0], 3, Rel.EQ, "scaler_f_shape", cls_name) + validator.check_int( + atom_b_shape[0], M, Rel.EQ, "atom_b_shape", cls_name) + validator.check_int( + bond_k_shape[0], M, Rel.EQ, "bond_k_shape", cls_name) + validator.check_int( + bond_r0_shape[0], M, Rel.EQ, "bond_r0_shape", cls_name) + + return bond_k_shape + def infer_dtype(self, uint_crd_f_dtype, scaler_f_type, atom_a_type, atom_b_type, bond_k_type, bond_r0_type): validator.check_tensor_dtype_valid('uint_crd_f_dtype', uint_crd_f_dtype, [mstype.uint32], self.name) validator.check_tensor_dtype_valid('scaler_f_type', scaler_f_type, [mstype.float32], self.name) @@ -125,12 +168,30 @@ class BondAtomEnergy(PrimitiveWithInfer): ``GPU`` Examples: """ + @prim_attr_register def __init__(self, bond_numbers): self.bond_numbers = bond_numbers self.init_prim_io_names(inputs=['uint_crd_f', 'scaler_f', 'atom_a', 'atom_b', 'bond_k', 'bond_r0'], outputs=['atom_ene']) self.add_prim_attr('bond_numbers', self.bond_numbers) + + 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 = uint_crd_f_shape[0] + M = atom_a_shape[0] + validator.check_int( + uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape", cls_name) + validator.check_int( + scaler_f_shape[0], 3, Rel.EQ, "scaler_f_shape", cls_name) + validator.check_int( + atom_b_shape[0], M, Rel.EQ, "atom_b_shape", cls_name) + validator.check_int( + bond_k_shape[0], M, Rel.EQ, "bond_k_shape", cls_name) + validator.check_int( + bond_r0_shape[0], M, Rel.EQ, "bond_r0_shape", cls_name) + return [N,] + def infer_dtype(self, uint_crd_f_dtype, scaler_f_type, atom_a_type, atom_b_type, bond_k_type, bond_r0_type): validator.check_tensor_dtype_valid('uint_crd_f_dtype', uint_crd_f_dtype, [mstype.uint32], self.name) validator.check_tensor_dtype_valid('scaler_f_type', scaler_f_type, [mstype.float32], self.name) @@ -167,13 +228,28 @@ class BondForceWithAtomEnergy(PrimitiveWithInfer): self.init_prim_io_names(inputs=['uint_crd_f', 'scaler_f', 'atom_a', 'atom_b', 'bond_k', 'bond_r0'], outputs=['frc_f', 'atom_e']) self.add_prim_attr('bond_numbers', self.bond_numbers) + + 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 = uint_crd_f_shape[0] + M = atom_a_shape[0] + validator.check_int( + uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape", cls_name) + validator.check_int( + scaler_f_shape[0], 3, Rel.EQ, "scaler_f_shape", cls_name) + validator.check_int( + atom_b_shape[0], M, Rel.EQ, "atom_b_shape", cls_name) + validator.check_int( + bond_k_shape[0], M, Rel.EQ, "bond_k_shape", cls_name) + validator.check_int( + bond_r0_shape[0], M, Rel.EQ, "bond_r0_shape", cls_name) + return uint_crd_f_shape, [N,] + def infer_dtype(self, uint_crd_f_dtype, scaler_f_type, atom_a_type, atom_b_type, bond_k_type, bond_r0_type): validator.check_tensor_dtype_valid('uint_crd_f_dtype', uint_crd_f_dtype, [mstype.uint32], self.name) validator.check_tensor_dtype_valid('scaler_f_type', scaler_f_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('atom_a_type', atom_a_type, [mstype.int32], self.name) validator.check_tensor_dtype_valid('atom_b_type', atom_b_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('bond_k_type', bond_k_type, [mstype.float32], self.name) validator.check_tensor_dtype_valid('bond_r0_type', bond_r0_type, [mstype.float32], self.name) return bond_r0_type, bond_r0_type @@ -213,17 +289,33 @@ class BondForceWithAtomVirial(PrimitiveWithInfer): self.init_prim_io_names(inputs=['uint_crd_f', 'scaler_f', 'atom_a', 'atom_b', 'bond_k', 'bond_r0'], outputs=['frc_f', 'atom_v']) self.add_prim_attr('bond_numbers', self.bond_numbers) + + 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 = uint_crd_f_shape[0] + M = atom_a_shape[0] + validator.check_int( + uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape", cls_name) + validator.check_int( + scaler_f_shape[0], 3, Rel.EQ, "scaler_f_shape", cls_name) + validator.check_int( + atom_b_shape[0], M, Rel.EQ, "atom_b_shape", cls_name) + validator.check_int( + bond_k_shape[0], M, Rel.EQ, "bond_k_shape", cls_name) + validator.check_int( + bond_r0_shape[0], M, Rel.EQ, "bond_r0_shape", cls_name) + return uint_crd_f_shape, [N,] + def infer_dtype(self, uint_crd_f_dtype, scaler_f_type, atom_a_type, atom_b_type, bond_k_type, bond_r0_type): validator.check_tensor_dtype_valid('uint_crd_f_dtype', uint_crd_f_dtype, [mstype.uint32], self.name) validator.check_tensor_dtype_valid('scaler_f_type', scaler_f_type, [mstype.float32], self.name) - validator.check_tensor_dtype_valid('atom_a_type', atom_a_type, [mstype.int32], self.name) validator.check_tensor_dtype_valid('atom_b_type', atom_b_type, [mstype.int32], self.name) - validator.check_tensor_dtype_valid('bond_k_type', bond_k_type, [mstype.float32], self.name) validator.check_tensor_dtype_valid('bond_r0_type', bond_r0_type, [mstype.float32], self.name) return bond_r0_type, bond_r0_type + class DihedralForce(PrimitiveWithInfer): """ DihedralForce: @@ -259,18 +351,18 @@ class DihedralForce(PrimitiveWithInfer): Inputs: - **uint_crd_f** (Tensor, uint32) - [N, 3], the unsigned int coordinates value of each atom. - - **scalar_f** (Tensor, float32) - [3, 1], the 3-D scale factor between + - **scalar_f** (Tensor, float32) - [3, ], the 3-D scale factor between the real space float coordinates and the unsigned int coordinates. - - **atom_a** (Tensor, int32) - [M, 1], the 1st atom index of each dihedral. - - **atom_b** (Tensor, int32) - [M, 1], the 2nd atom index of each dihedral. - - **atom_c** (Tensor, int32) - [M, 1], the 3rd atom index of each dihedral. - - **atom_d** (Tensor, int32) - [M, 1], the 4th atom index of each dihedral. + - **atom_a** (Tensor, int32) - [M, ], the 1st atom index of each dihedral. + - **atom_b** (Tensor, int32) - [M, ], the 2nd atom index of each dihedral. + - **atom_c** (Tensor, int32) - [M, ], the 3rd atom index of each dihedral. + - **atom_d** (Tensor, int32) - [M, ], the 4th atom index of each dihedral. 4 atoms are connected in the form a-b-c-d. - - **ipn** (Tensor, int32) - [M, 1], the period of dihedral angle of each dihedral. - - **pk** (Tensor, float32) - [M, 1], the force constant of each dihedral. - - **gamc** (Tensor, float32) - [M, 1], k*cos(phi_0) of each dihedral. - - **gams** (Tensor, float32) - [M, 1], k*sin(phi_0) of each dihedral. - - **pn** (Tensor, float32) - [M, 1], the floating point form of ipn. + - **ipn** (Tensor, int32) - [M, ], the period of dihedral angle of each dihedral. + - **pk** (Tensor, float32) - [M, ], the force constant of each dihedral. + - **gamc** (Tensor, float32) - [M, ], k*cos(phi_0) of each dihedral. + - **gams** (Tensor, float32) - [M, ], k*sin(phi_0) of each dihedral. + - **pn** (Tensor, float32) - [M, ], the floating point form of ipn. Outputs: - **frc_f** (Tensor, float32) - [N, 3], the force felt by each atom. @@ -289,6 +381,29 @@ class DihedralForce(PrimitiveWithInfer): outputs=['frc_f']) self.add_prim_attr('dihedral_numbers', self.dihedral_numbers) + def infer_shape(self, uint_crd_f_shape, scaler_f_shape, atom_a_shape, atom_b_shape, atom_c_shape, atom_d_shape, + ipn_shape, pk_shape, gamc_shape, gams_shape, pn_shape): + cls_name = self.name + M = atom_a_shape[0] + validator.check_int( + uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape", cls_name) + validator.check_int( + scaler_f_shape[0], 3, Rel.EQ, "scaler_f_shape", cls_name) + validator.check_int( + atom_a_shape[0], M, Rel.EQ, "atom_a_shape", cls_name) + validator.check_int( + atom_b_shape[0], M, Rel.EQ, "atom_b_shape", cls_name) + validator.check_int( + atom_c_shape[0], M, Rel.EQ, "atom_c_shape", cls_name) + validator.check_int( + atom_d_shape[0], M, Rel.EQ, "atom_d_shape", cls_name) + validator.check_int(ipn_shape[0], M, Rel.EQ, "ipn_shape", cls_name) + validator.check_int(pk_shape[0], M, Rel.EQ, "pk_shape", cls_name) + validator.check_int(gamc_shape[0], M, Rel.EQ, "gamc_shape", cls_name) + validator.check_int(gams_shape[0], M, Rel.EQ, "gams_shape", cls_name) + validator.check_int(pn_shape[0], M, Rel.EQ, "pn_shape", cls_name) + return uint_crd_f_shape + def infer_dtype(self, uint_crd_f_dtype, scaler_f_type, atom_a_type, atom_b_type, atom_c_type, atom_d_type, ipn_type, pk_type, gamc_type, gams_type, pn_type): validator.check_tensor_dtype_valid('uint_crd_f_dtype', uint_crd_f_dtype, [mstype.uint32], self.name) @@ -302,7 +417,6 @@ class DihedralForce(PrimitiveWithInfer): validator.check_tensor_dtype_valid('gamc_type', gamc_type, [mstype.float32], self.name) validator.check_tensor_dtype_valid('gams_type', gams_type, [mstype.float32], self.name) validator.check_tensor_dtype_valid('pn_type', pn_type, [mstype.float32], self.name) - return pn_type @@ -321,7 +435,7 @@ class DihedralEnergy(PrimitiveWithInfer): Same as operator DihedralForce(). Outputs: - - **ene** (Tensor, float32) - [M, 1], the potential energy for each + - **ene** (Tensor, float32) - [M, ], the potential energy for each dihedral term. Supported Platforms: @@ -338,6 +452,29 @@ class DihedralEnergy(PrimitiveWithInfer): outputs=['ene']) self.add_prim_attr('dihedral_numbers', self.dihedral_numbers) + def infer_shape(self, uint_crd_f_shape, scaler_f_shape, atom_a_shape, atom_b_shape, atom_c_shape, atom_d_shape, + ipn_shape, pk_shape, gamc_shape, gams_shape, pn_shape): + cls_name = self.name + M = atom_a_shape[0] + validator.check_int( + uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape", cls_name) + validator.check_int( + scaler_f_shape[0], 3, Rel.EQ, "scaler_f_shape", cls_name) + validator.check_int( + atom_a_shape[0], M, Rel.EQ, "atom_a_shape", cls_name) + validator.check_int( + atom_b_shape[0], M, Rel.EQ, "atom_b_shape", cls_name) + validator.check_int( + atom_c_shape[0], M, Rel.EQ, "atom_c_shape", cls_name) + validator.check_int( + atom_d_shape[0], M, Rel.EQ, "atom_d_shape", cls_name) + validator.check_int(ipn_shape[0], M, Rel.EQ, "ipn_shape", cls_name) + validator.check_int(pk_shape[0], M, Rel.EQ, "pk_shape", cls_name) + validator.check_int(gamc_shape[0], M, Rel.EQ, "gamc_shape", cls_name) + validator.check_int(gams_shape[0], M, Rel.EQ, "gams_shape", cls_name) + validator.check_int(pn_shape[0], M, Rel.EQ, "pn_shape", cls_name) + return [M,] + def infer_dtype(self, uint_crd_f_dtype, scaler_f_type, atom_a_type, atom_b_type, atom_c_type, atom_d_type, ipn_type, pk_type, gamc_type, gams_type, pn_type): validator.check_tensor_dtype_valid('uint_crd_f_dtype', uint_crd_f_dtype, [mstype.uint32], self.name) @@ -351,7 +488,6 @@ class DihedralEnergy(PrimitiveWithInfer): validator.check_tensor_dtype_valid('gamc_type', gamc_type, [mstype.float32], self.name) validator.check_tensor_dtype_valid('gams_type', gams_type, [mstype.float32], self.name) validator.check_tensor_dtype_valid('pn_type', pn_type, [mstype.float32], self.name) - return pn_type @@ -368,7 +504,7 @@ class DihedralAtomEnergy(PrimitiveWithInfer): Same as operator DihedralEnergy(). Outputs: - - **ene** (Tensor, float32) - [N, 1], the accumulated potential + - **ene** (Tensor, float32) - [N, ], the accumulated potential energy for each atom. Supported Platforms: @@ -385,6 +521,30 @@ class DihedralAtomEnergy(PrimitiveWithInfer): outputs=['ene']) self.add_prim_attr('dihedral_numbers', self.dihedral_numbers) + def infer_shape(self, uint_crd_f_shape, scaler_f_shape, atom_a_shape, atom_b_shape, atom_c_shape, atom_d_shape, + ipn_shape, pk_shape, gamc_shape, gams_shape, pn_shape): + cls_name = self.name + N = uint_crd_f_shape[0] + M = atom_a_shape[0] + validator.check_int( + uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape", cls_name) + validator.check_int( + scaler_f_shape[0], 3, Rel.EQ, "scaler_f_shape", cls_name) + validator.check_int( + atom_a_shape[0], M, Rel.EQ, "atom_a_shape", cls_name) + validator.check_int( + atom_b_shape[0], M, Rel.EQ, "atom_b_shape", cls_name) + validator.check_int( + atom_c_shape[0], M, Rel.EQ, "atom_c_shape", cls_name) + validator.check_int( + atom_d_shape[0], M, Rel.EQ, "atom_d_shape", cls_name) + validator.check_int(ipn_shape[0], M, Rel.EQ, "ipn_shape", cls_name) + validator.check_int(pk_shape[0], M, Rel.EQ, "pk_shape", cls_name) + validator.check_int(gamc_shape[0], M, Rel.EQ, "gamc_shape", cls_name) + validator.check_int(gams_shape[0], M, Rel.EQ, "gams_shape", cls_name) + validator.check_int(pn_shape[0], M, Rel.EQ, "pn_shape", cls_name) + return [N,] + def infer_dtype(self, uint_crd_f_dtype, scaler_f_type, atom_a_type, atom_b_type, atom_c_type, atom_d_type, ipn_type, pk_type, gamc_type, gams_type, pn_type): validator.check_tensor_dtype_valid('uint_crd_f_dtype', uint_crd_f_dtype, [mstype.uint32], self.name) @@ -398,7 +558,6 @@ class DihedralAtomEnergy(PrimitiveWithInfer): validator.check_tensor_dtype_valid('gamc_type', gamc_type, [mstype.float32], self.name) validator.check_tensor_dtype_valid('gams_type', gams_type, [mstype.float32], self.name) validator.check_tensor_dtype_valid('pn_type', pn_type, [mstype.float32], self.name) - return pn_type @@ -415,7 +574,7 @@ class DihedralForceWithAtomEnergy(PrimitiveWithInfer): Outputs: - **frc_f** (Tensor, float32) - [N, 3], same as operator DihedralForce(). - - **ene** (Tensor, float32) - [N, 1], same as operator DihedralAtomEnergy(). + - **ene** (Tensor, float32) - [N, ], same as operator DihedralAtomEnergy(). Supported Platforms: ``GPU`` @@ -431,6 +590,30 @@ class DihedralForceWithAtomEnergy(PrimitiveWithInfer): outputs=['frc_f', 'ene']) self.add_prim_attr('dihedral_numbers', self.dihedral_numbers) + def infer_shape(self, uint_crd_f_shape, scaler_f_shape, atom_a_shape, atom_b_shape, atom_c_shape, atom_d_shape, + ipn_shape, pk_shape, gamc_shape, gams_shape, pn_shape): + cls_name = self.name + N = uint_crd_f_shape[0] + M = atom_a_shape[0] + validator.check_int( + uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape", cls_name) + validator.check_int( + scaler_f_shape[0], 3, Rel.EQ, "scaler_f_shape", cls_name) + validator.check_int( + atom_a_shape[0], M, Rel.EQ, "atom_a_shape", cls_name) + validator.check_int( + atom_b_shape[0], M, Rel.EQ, "atom_b_shape", cls_name) + validator.check_int( + atom_c_shape[0], M, Rel.EQ, "atom_c_shape", cls_name) + validator.check_int( + atom_d_shape[0], M, Rel.EQ, "atom_d_shape", cls_name) + validator.check_int(ipn_shape[0], M, Rel.EQ, "ipn_shape", cls_name) + validator.check_int(pk_shape[0], M, Rel.EQ, "pk_shape", cls_name) + validator.check_int(gamc_shape[0], M, Rel.EQ, "gamc_shape", cls_name) + validator.check_int(gams_shape[0], M, Rel.EQ, "gams_shape", cls_name) + validator.check_int(pn_shape[0], M, Rel.EQ, "pn_shape", cls_name) + return uint_crd_f_shape, [N,] + def infer_dtype(self, uint_crd_f_dtype, scaler_f_type, atom_a_type, atom_b_type, atom_c_type, atom_d_type, ipn_type, pk_type, gamc_type, gams_type, pn_type): validator.check_tensor_dtype_valid('uint_crd_f_dtype', uint_crd_f_dtype, [mstype.uint32], self.name) @@ -444,7 +627,6 @@ class DihedralForceWithAtomEnergy(PrimitiveWithInfer): validator.check_tensor_dtype_valid('gamc_type', gamc_type, [mstype.float32], self.name) validator.check_tensor_dtype_valid('gams_type', gams_type, [mstype.float32], self.name) validator.check_tensor_dtype_valid('pn_type', pn_type, [mstype.float32], self.name) - return pn_type, pn_type @@ -470,14 +652,14 @@ class AngleForce(PrimitiveWithInfer): Inputs: - **uint_crd_f** (Tensor, uint32) - [N, 3], the unsigned int coordinate value of each atom. - - **scaler_f** (Tensor, float32) - [3, 1], the 3-D scale factor between + - **scaler_f** (Tensor, float32) - [3, ], the 3-D scale factor between the real space float coordinates and the unsigned int coordinates. - - **atom_a** (Tensor, int32) - [M, 1], the 1st atom index of each angle. - - **atom_b** (Tensor, int32) - [M, 1], the 2nd and the central atom index + - **atom_a** (Tensor, int32) - [M, ], the 1st atom index of each angle. + - **atom_b** (Tensor, int32) - [M, ], the 2nd and the central atom index of each angle. - - **atom_c** (Tensor, int32) - [M, 1], the 3rd atom index of each angle. - - **angle_k** (Tensor, float32) - [M, 1], the force constant for each angle. - - **angle_theta0** (Tensor, float32) - [M, 1], the equilibrium position value + - **atom_c** (Tensor, int32) - [M, ], the 3rd atom index of each angle. + - **angle_k** (Tensor, float32) - [M, ], the force constant for each angle. + - **angle_theta0** (Tensor, float32) - [M, ], the equilibrium position value for each angle. Outputs: @@ -497,6 +679,26 @@ class AngleForce(PrimitiveWithInfer): outputs=['frc_f']) self.add_prim_attr('angle_numbers', self.angle_numbers) + def infer_shape(self, uint_crd_f_shape, scaler_f_shape, atom_a_shape, atom_b_shape, atom_c_shape, angle_k_shape, + angle_theta0_shape): + cls_name = self.name + M = atom_a_shape[0] + validator.check_int( + uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape", cls_name) + validator.check_int( + scaler_f_shape[0], 3, Rel.EQ, "scaler_f_shape", cls_name) + validator.check_int( + atom_a_shape[0], M, Rel.EQ, "atom_a_shape", cls_name) + validator.check_int( + atom_b_shape[0], M, Rel.EQ, "atom_b_shape", cls_name) + validator.check_int( + atom_c_shape[0], M, Rel.EQ, "atom_c_shape", cls_name) + validator.check_int( + angle_k_shape[0], M, Rel.EQ, "angle_k_shape", cls_name) + validator.check_int( + angle_theta0_shape[0], M, Rel.EQ, "angle_theta0_shape", cls_name) + return uint_crd_f_shape + def infer_dtype(self, uint_crd_f_dtype, scaler_f_type, atom_a_type, atom_b_type, atom_c_type, angle_k_type, angle_theta0_type): validator.check_tensor_dtype_valid('uint_crd_f_dtype', uint_crd_f_dtype, [mstype.uint32], self.name) @@ -526,7 +728,7 @@ class AngleEnergy(PrimitiveWithInfer): Same as operator AngleForce(). Outputs: - - **ene** (Tensor, float32) - [M, 1], the potential energy for + - **ene** (Tensor, float32) - [M, ], the potential energy for each angle term. Supported Platforms: @@ -543,6 +745,26 @@ class AngleEnergy(PrimitiveWithInfer): outputs=['ene']) self.add_prim_attr('angle_numbers', self.angle_numbers) + def infer_shape(self, uint_crd_f_shape, scaler_f_shape, atom_a_shape, atom_b_shape, atom_c_shape, angle_k_shape, + angle_theta0_shape): + cls_name = self.name + M = atom_a_shape[0] + validator.check_int( + uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape", cls_name) + validator.check_int( + scaler_f_shape[0], 3, Rel.EQ, "scaler_f_shape", cls_name) + validator.check_int( + atom_a_shape[0], M, Rel.EQ, "atom_a_shape", cls_name) + validator.check_int( + atom_b_shape[0], M, Rel.EQ, "atom_b_shape", cls_name) + validator.check_int( + atom_c_shape[0], M, Rel.EQ, "atom_c_shape", cls_name) + validator.check_int( + angle_k_shape[0], M, Rel.EQ, "angle_k_shape", cls_name) + validator.check_int( + angle_theta0_shape[0], M, Rel.EQ, "angle_theta0_shape", cls_name) + return [M,] + def infer_dtype(self, uint_crd_f_dtype, scaler_f_type, atom_a_type, atom_b_type, atom_c_type, angle_k_type, angle_theta0_type): validator.check_tensor_dtype_valid('uint_crd_f_dtype', uint_crd_f_dtype, [mstype.uint32], self.name) @@ -568,7 +790,7 @@ class AngleAtomEnergy(PrimitiveWithInfer): Same as operator AngleForce(). Outputs: - - **ene** (Tensor, float32) - [N, 1], the accumulated potential energy + - **ene** (Tensor, float32) - [N, ], the accumulated potential energy for each atom. Supported Platforms: @@ -585,6 +807,27 @@ class AngleAtomEnergy(PrimitiveWithInfer): outputs=['ene']) self.add_prim_attr('angle_numbers', self.angle_numbers) + def infer_shape(self, uint_crd_f_shape, scaler_f_shape, atom_a_shape, atom_b_shape, atom_c_shape, angle_k_shape, + angle_theta0_shape): + cls_name = self.name + N = uint_crd_f_shape[0] + M = atom_a_shape[0] + validator.check_int( + uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape", cls_name) + validator.check_int( + scaler_f_shape[0], 3, Rel.EQ, "scaler_f_shape", cls_name) + validator.check_int( + atom_a_shape[0], M, Rel.EQ, "atom_a_shape", cls_name) + validator.check_int( + atom_b_shape[0], M, Rel.EQ, "atom_b_shape", cls_name) + validator.check_int( + atom_c_shape[0], M, Rel.EQ, "atom_c_shape", cls_name) + validator.check_int( + angle_k_shape[0], M, Rel.EQ, "angle_k_shape", cls_name) + validator.check_int( + angle_theta0_shape[0], M, Rel.EQ, "angle_theta0_shape", cls_name) + return [N,] + def infer_dtype(self, uint_crd_f_dtype, scaler_f_type, atom_a_type, atom_b_type, atom_c_type, angle_k_type, angle_theta0_type): validator.check_tensor_dtype_valid('uint_crd_f_dtype', uint_crd_f_dtype, [mstype.uint32], self.name) @@ -610,7 +853,7 @@ class AngleForceWithAtomEnergy(PrimitiveWithInfer): Outputs: - **frc_f** (Tensor, float32) - [N, 3], same as operator AngleForce(). - - **ene** (Tensor, float) - [N, 1], same as operator AngleAtomEnergy(). + - **ene** (Tensor, float) - [N, ], same as operator AngleAtomEnergy(). Supported Platforms: ``GPU`` @@ -626,6 +869,27 @@ class AngleForceWithAtomEnergy(PrimitiveWithInfer): outputs=['frc_f', 'ene']) self.add_prim_attr('angle_numbers', self.angle_numbers) + def infer_shape(self, uint_crd_f_shape, scaler_f_shape, atom_a_shape, atom_b_shape, atom_c_shape, angle_k_shape, + angle_theta0_shape): + cls_name = self.name + N = uint_crd_f_shape[0] + M = atom_a_shape[0] + validator.check_int( + uint_crd_f_shape[1], 3, Rel.EQ, "uint_crd_f_shape", cls_name) + validator.check_int( + scaler_f_shape[0], 3, Rel.EQ, "scaler_f_shape", cls_name) + validator.check_int( + atom_a_shape[0], M, Rel.EQ, "atom_a_shape", cls_name) + validator.check_int( + atom_b_shape[0], M, Rel.EQ, "atom_b_shape", cls_name) + validator.check_int( + atom_c_shape[0], M, Rel.EQ, "atom_c_shape", cls_name) + validator.check_int( + angle_k_shape[0], M, Rel.EQ, "angle_k_shape", cls_name) + validator.check_int( + angle_theta0_shape[0], M, Rel.EQ, "angle_theta0_shape", cls_name) + return uint_crd_f_shape, [N,] + def infer_dtype(self, uint_crd_f_dtype, scaler_f_type, atom_a_type, atom_b_type, atom_c_type, angle_k_type, angle_theta0_type): validator.check_tensor_dtype_valid('uint_crd_f_dtype', uint_crd_f_dtype, [mstype.uint32], self.name)