From fac79a10a62fdeaceaa53e8e19d1015bd85df842 Mon Sep 17 00:00:00 2001 From: huangbingjian Date: Mon, 27 Sep 2021 19:54:37 +0800 Subject: [PATCH] add side-effect mark to sponge operators --- mindspore/ops/operations/sponge_ops.py | 1 + mindspore/ops/operations/sponge_update_ops.py | 5 +++++ 2 files changed, 6 insertions(+) diff --git a/mindspore/ops/operations/sponge_ops.py b/mindspore/ops/operations/sponge_ops.py index 84c55e7c1bc..7b0f0b8bff0 100644 --- a/mindspore/ops/operations/sponge_ops.py +++ b/mindspore/ops/operations/sponge_ops.py @@ -2791,6 +2791,7 @@ class MDIterationLeapFrogLiujian(PrimitiveWithInfer): 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 diff --git a/mindspore/ops/operations/sponge_update_ops.py b/mindspore/ops/operations/sponge_update_ops.py index 970e2e60d59..e39b672b7ec 100644 --- a/mindspore/ops/operations/sponge_update_ops.py +++ b/mindspore/ops/operations/sponge_update_ops.py @@ -312,6 +312,7 @@ class RefreshCrdVel(PrimitiveWithInfer): 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 @@ -625,6 +626,7 @@ class MDIterationLeapFrogLiujianWithMaxVel(PrimitiveWithInfer): 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 @@ -1085,6 +1087,7 @@ class MDIterationLeapFrog(PrimitiveWithInfer): 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 @@ -1162,6 +1165,7 @@ class MDIterationLeapFrogWithMaxVel(PrimitiveWithInfer): 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 @@ -1227,6 +1231,7 @@ class MDIterationGradientDescent(PrimitiveWithInfer): 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