diff --git a/source/op/tf/custom_op.h b/source/op/tf/custom_op.h index 538bfeb2cc..a10b4f4d1b 100644 --- a/source/op/tf/custom_op.h +++ b/source/op/tf/custom_op.h @@ -1,6 +1,7 @@ // SPDX-License-Identifier: LGPL-3.0-or-later #pragma once #include +#include #include #include #include @@ -51,6 +52,48 @@ inline Status InvalidArgument(Args&&... args) { return tensorflow::errors::InvalidArgument(std::forward(args)...); #endif } + +/** + * @brief Derive a dense tensor's per-atom width without truncating division. + * + * Several low-level TensorFlow ops flatten atom and feature dimensions into a + * single axis. Validate the flattened width before dividing by `nloc`; raw + * CPU/GPU kernels cannot safely consume a leftover partial atom row. + * + * @param per_atom_width Receives the validated feature width for one atom. + * @param shape Rank-two tensor shape whose second dimension is flattened. + * @param nloc Number of local atoms encoded in the flattened dimension. + * @param tensor_name Human-readable input name used in validation errors. + * @return An OK status, or InvalidArgument when the width is incompatible. + */ +inline Status GetPerAtomWidth(int* per_atom_width, + const TensorShape& shape, + const int nloc, + const char* tensor_name) { + const int64_t flattened_width = shape.dim_size(1); + if (nloc < 0) { + return InvalidArgument("number of local atoms should be non-negative"); + } + if (nloc == 0) { + if (flattened_width != 0) { + return InvalidArgument(tensor_name, + " width should be zero when nloc is zero"); + } + *per_atom_width = 0; + return Status(); + } + if (flattened_width % nloc != 0) { + return InvalidArgument(tensor_name, " width ", flattened_width, + " should be divisible by nloc ", nloc); + } + const int64_t width = flattened_width / nloc; + if (width > std::numeric_limits::max()) { + return InvalidArgument(tensor_name, + " width per atom exceeds the supported int range"); + } + *per_atom_width = static_cast(width); + return Status(); +} } // namespace tf_compat } // namespace deepmd diff --git a/source/op/tf/prod_force_grad_multi_device.cc b/source/op/tf/prod_force_grad_multi_device.cc index 6678c1b628..39a26cde68 100644 --- a/source/op/tf/prod_force_grad_multi_device.cc +++ b/source/op/tf/prod_force_grad_multi_device.cc @@ -73,8 +73,12 @@ class ProdForceSeAGradOp : public OpKernel { int nframes = net_deriv_tensor.shape().dim_size(0); int nloc = natoms(0); - int ndescrpt = nloc > 0 ? net_deriv_tensor.shape().dim_size(1) / nloc : 0; - int nnei = nloc > 0 ? nlist_tensor.shape().dim_size(1) / nloc : 0; + int ndescrpt; + int nnei; + OP_REQUIRES_OK(context, deepmd::tf_compat::GetPerAtomWidth( + &ndescrpt, net_deriv_shape, nloc, "net deriv")); + OP_REQUIRES_OK(context, deepmd::tf_compat::GetPerAtomWidth( + &nnei, nlist_shape, nloc, "nlist")); // check the sizes OP_REQUIRES( @@ -97,6 +101,9 @@ class ProdForceSeAGradOp : public OpKernel { OP_REQUIRES( context, (nnei == n_a_sel + n_r_sel), deepmd::tf_compat::InvalidArgument("number of neighbors should match")); + OP_REQUIRES(context, (static_cast(nnei) * 4 == ndescrpt), + deepmd::tf_compat::InvalidArgument( + "descriptor width should be four times neighbor width")); // Create an output tensor TensorShape grad_net_shape; @@ -194,8 +201,12 @@ class ProdForceSeRGradOp : public OpKernel { int nframes = net_deriv_tensor.shape().dim_size(0); int nloc = natoms(0); - int ndescrpt = nloc > 0 ? net_deriv_tensor.shape().dim_size(1) / nloc : 0; - int nnei = nloc > 0 ? nlist_tensor.shape().dim_size(1) / nloc : 0; + int ndescrpt; + int nnei; + OP_REQUIRES_OK(context, deepmd::tf_compat::GetPerAtomWidth( + &ndescrpt, net_deriv_shape, nloc, "net deriv")); + OP_REQUIRES_OK(context, deepmd::tf_compat::GetPerAtomWidth( + &nnei, nlist_shape, nloc, "nlist")); // check the sizes OP_REQUIRES( @@ -215,6 +226,9 @@ class ProdForceSeRGradOp : public OpKernel { (int_64(nloc) * ndescrpt * 3 == in_deriv_shape.dim_size(1)), deepmd::tf_compat::InvalidArgument( "number of descriptors should match")); + OP_REQUIRES(context, (nnei == ndescrpt), + deepmd::tf_compat::InvalidArgument( + "descriptor width should equal neighbor width")); // Create an output tensor TensorShape grad_net_shape; diff --git a/source/op/tf/prod_force_multi_device.cc b/source/op/tf/prod_force_multi_device.cc index ae66f13503..3dcce46821 100644 --- a/source/op/tf/prod_force_multi_device.cc +++ b/source/op/tf/prod_force_multi_device.cc @@ -92,8 +92,13 @@ class ProdForceSeAOp : public OpKernel { int nloc = natoms[0]; int nall = natoms[1]; int nframes = net_deriv_tensor.shape().dim_size(0); - int ndescrpt = nloc > 0 ? net_deriv_tensor.shape().dim_size(1) / nloc : 0; - int nnei = nloc > 0 ? nlist_tensor.shape().dim_size(1) / nloc : 0; + int ndescrpt; + int nnei; + OP_REQUIRES_OK(context, + deepmd::tf_compat::GetPerAtomWidth( + &ndescrpt, net_deriv_tensor.shape(), nloc, "net deriv")); + OP_REQUIRES_OK(context, deepmd::tf_compat::GetPerAtomWidth( + &nnei, nlist_tensor.shape(), nloc, "nlist")); // check the sizes OP_REQUIRES( context, (nframes == in_deriv_tensor.shape().dim_size(0)), @@ -106,6 +111,9 @@ class ProdForceSeAOp : public OpKernel { (int_64(nloc) * ndescrpt * 3 == in_deriv_tensor.shape().dim_size(1)), deepmd::tf_compat::InvalidArgument( "number of descriptors should match")); + OP_REQUIRES(context, (static_cast(nnei) * 4 == ndescrpt), + deepmd::tf_compat::InvalidArgument( + "descriptor width should be four times neighbor width")); // Create an output tensor TensorShape force_shape; force_shape.AddDim(nframes); @@ -199,8 +207,13 @@ class ProdForceSeROp : public OpKernel { int nloc = natoms[0]; int nall = natoms[1]; int nframes = net_deriv_tensor.shape().dim_size(0); - int ndescrpt = nloc > 0 ? net_deriv_tensor.shape().dim_size(1) / nloc : 0; - int nnei = nloc > 0 ? nlist_tensor.shape().dim_size(1) / nloc : 0; + int ndescrpt; + int nnei; + OP_REQUIRES_OK(context, + deepmd::tf_compat::GetPerAtomWidth( + &ndescrpt, net_deriv_tensor.shape(), nloc, "net deriv")); + OP_REQUIRES_OK(context, deepmd::tf_compat::GetPerAtomWidth( + &nnei, nlist_tensor.shape(), nloc, "nlist")); // check the sizes OP_REQUIRES( context, (nframes == in_deriv_tensor.shape().dim_size(0)), @@ -213,6 +226,9 @@ class ProdForceSeROp : public OpKernel { in_deriv_tensor.shape().dim_size(1)), deepmd::tf_compat::InvalidArgument( "number of descriptors should match")); + OP_REQUIRES(context, (nnei == ndescrpt), + deepmd::tf_compat::InvalidArgument( + "descriptor width should equal neighbor width")); // Create an output tensor TensorShape force_shape; force_shape.AddDim(nframes); diff --git a/source/op/tf/prod_virial_grad_multi_device.cc b/source/op/tf/prod_virial_grad_multi_device.cc index fb011cabbd..18c00cc0de 100644 --- a/source/op/tf/prod_virial_grad_multi_device.cc +++ b/source/op/tf/prod_virial_grad_multi_device.cc @@ -79,8 +79,12 @@ class ProdVirialSeAGradOp : public OpKernel { int nframes = net_deriv_tensor.shape().dim_size(0); int nloc = natoms(0); - int ndescrpt = nloc > 0 ? net_deriv_tensor.shape().dim_size(1) / nloc : 0; - int nnei = nloc > 0 ? nlist_tensor.shape().dim_size(1) / nloc : 0; + int ndescrpt; + int nnei; + OP_REQUIRES_OK(context, deepmd::tf_compat::GetPerAtomWidth( + &ndescrpt, net_deriv_shape, nloc, "net deriv")); + OP_REQUIRES_OK(context, deepmd::tf_compat::GetPerAtomWidth( + &nnei, nlist_shape, nloc, "nlist")); // check the sizes OP_REQUIRES( @@ -109,6 +113,9 @@ class ProdVirialSeAGradOp : public OpKernel { OP_REQUIRES( context, (nnei == n_a_sel + n_r_sel), deepmd::tf_compat::InvalidArgument("number of neighbors should match")); + OP_REQUIRES(context, (static_cast(nnei) * 4 == ndescrpt), + deepmd::tf_compat::InvalidArgument( + "descriptor width should be four times neighbor width")); // Create an output tensor TensorShape grad_net_shape; @@ -223,8 +230,12 @@ class ProdVirialSeRGradOp : public OpKernel { int nframes = net_deriv_tensor.shape().dim_size(0); int nloc = natoms(0); - int ndescrpt = nloc > 0 ? net_deriv_tensor.shape().dim_size(1) / nloc : 0; - int nnei = nloc > 0 ? nlist_tensor.shape().dim_size(1) / nloc : 0; + int ndescrpt; + int nnei; + OP_REQUIRES_OK(context, deepmd::tf_compat::GetPerAtomWidth( + &ndescrpt, net_deriv_shape, nloc, "net deriv")); + OP_REQUIRES_OK(context, deepmd::tf_compat::GetPerAtomWidth( + &nnei, nlist_shape, nloc, "nlist")); // check the sizes OP_REQUIRES( @@ -250,6 +261,9 @@ class ProdVirialSeRGradOp : public OpKernel { OP_REQUIRES( context, (int_64(nloc) * nnei * 3 == rij_shape.dim_size(1)), deepmd::tf_compat::InvalidArgument("dim of rij should be nnei * 3")); + OP_REQUIRES(context, (nnei == ndescrpt), + deepmd::tf_compat::InvalidArgument( + "descriptor width should equal neighbor width")); // Create an output tensor TensorShape grad_net_shape; diff --git a/source/tests/tf/test_multi_device_shape_validation.py b/source/tests/tf/test_multi_device_shape_validation.py new file mode 100644 index 0000000000..5a7f0bf2be --- /dev/null +++ b/source/tests/tf/test_multi_device_shape_validation.py @@ -0,0 +1,217 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Regression tests for flattened TensorFlow custom-op input dimensions.""" + +from deepmd.tf.env import ( + GLOBAL_TF_FLOAT_PRECISION, + op_grads_module, + op_module, + tf, +) + + +class TestMultiDeviceShapeValidation(tf.test.TestCase): + """Ensure malformed flattened widths fail before native kernel dispatch.""" + + def setUp(self) -> None: + self.sess = self.cached_session().__enter__() + self.nloc = 2 + self.nnei = 1 + self.ndescrpt = 4 + self.natoms = tf.constant([self.nloc, self.nloc, 1], dtype=tf.int32) + + def _floats(self, width: int): + """Create one frame of flattened floating-point custom-op input.""" + return tf.zeros([1, width], dtype=GLOBAL_TF_FLOAT_PRECISION) + + def _nlist(self, width: int): + """Create one frame of flattened neighbor indices.""" + return tf.zeros([1, width], dtype=tf.int32) + + def test_negative_nloc_is_rejected(self) -> None: + natoms = tf.constant([-1, 0, 0], dtype=tf.int32) + + with self.assertRaisesRegex( + tf.errors.InvalidArgumentError, + r"number of local atoms should be non-negative", + ): + self.sess.run( + op_module.prod_force_se_a( + self._floats(0), + self._floats(0), + self._nlist(0), + natoms, + n_a_sel=0, + n_r_sel=0, + ) + ) + + def test_zero_nloc_rejects_nonempty_flattened_width(self) -> None: + natoms = tf.constant([0, 0, 0], dtype=tf.int32) + + with self.assertRaisesRegex( + tf.errors.InvalidArgumentError, + r"net deriv width should be zero when nloc is zero", + ): + self.sess.run( + op_module.prod_force_se_a( + self._floats(1), + self._floats(0), + self._nlist(0), + natoms, + n_a_sel=0, + n_r_sel=0, + ) + ) + + def test_zero_nloc_accepts_empty_flattened_widths(self) -> None: + natoms = tf.constant([0, 0, 0], dtype=tf.int32) + + result = self.sess.run( + op_module.prod_force_se_a( + self._floats(0), + self._floats(0), + self._nlist(0), + natoms, + n_a_sel=0, + n_r_sel=0, + ) + ) + + self.assertEqual(result.shape, (1, 0)) + + def test_prod_force_rejects_partial_net_deriv_atom(self) -> None: + # The old integer division truncated 9 / 2 to four descriptors and + # allowed the extra value to survive until raw pointer dispatch. + with self.assertRaisesRegex( + tf.errors.InvalidArgumentError, + r"net deriv width 9 should be divisible by nloc 2", + ): + self.sess.run( + op_module.prod_force_se_a( + self._floats(self.nloc * self.ndescrpt + 1), + self._floats(self.nloc * self.ndescrpt * 3), + self._nlist(self.nloc * self.nnei), + self.natoms, + n_a_sel=self.nnei, + n_r_sel=0, + ) + ) + + def test_prod_force_rejects_in_deriv_width_mismatch(self) -> None: + with self.assertRaisesRegex( + tf.errors.InvalidArgumentError, r"number of descriptors should match" + ): + self.sess.run( + op_module.prod_force_se_a( + self._floats(self.nloc * self.ndescrpt), + self._floats(self.nloc * self.ndescrpt * 3 - 1), + self._nlist(self.nloc * self.nnei), + self.natoms, + n_a_sel=self.nnei, + n_r_sel=0, + ) + ) + + def test_prod_force_r_rejects_descriptor_stride_mismatch(self) -> None: + with self.assertRaisesRegex( + tf.errors.InvalidArgumentError, + r"descriptor width should equal neighbor width", + ): + self.sess.run( + op_module.prod_force_se_r( + self._floats(self.nloc * (self.nnei + 1)), + self._floats(self.nloc * (self.nnei + 1) * 3), + self._nlist(self.nloc * self.nnei), + self.natoms, + ) + ) + + def test_prod_force_grad_rejects_partial_nlist_atom(self) -> None: + # Fixed-width placeholders in the original tests rejected this feed + # before the custom op ran, leaving its release-build checks untested. + with self.assertRaisesRegex( + tf.errors.InvalidArgumentError, + r"nlist width 3 should be divisible by nloc 2", + ): + self.sess.run( + op_grads_module.prod_force_se_a_grad( + self._floats(self.nloc * 3), + self._floats(self.nloc * self.ndescrpt), + self._floats(self.nloc * self.ndescrpt * 3), + self._nlist(self.nloc * self.nnei + 1), + self.natoms, + n_a_sel=self.nnei, + n_r_sel=0, + ) + ) + + def test_prod_force_r_grad_rejects_partial_nlist_atom(self) -> None: + with self.assertRaisesRegex( + tf.errors.InvalidArgumentError, + r"nlist width 3 should be divisible by nloc 2", + ): + self.sess.run( + op_grads_module.prod_force_se_r_grad( + self._floats(self.nloc * 3), + self._floats(self.nloc * self.nnei), + self._floats(self.nloc * self.nnei * 3), + self._nlist(self.nloc * self.nnei + 1), + self.natoms, + ) + ) + + def test_prod_virial_grad_rejects_descriptor_stride_mismatch(self) -> None: + mismatched_ndescrpt = self.ndescrpt * 2 + with self.assertRaisesRegex( + tf.errors.InvalidArgumentError, + r"descriptor width should be four times neighbor width", + ): + self.sess.run( + op_grads_module.prod_virial_se_a_grad( + self._floats(9), + self._floats(self.nloc * mismatched_ndescrpt), + self._floats(self.nloc * mismatched_ndescrpt * 3), + self._floats(self.nloc * self.nnei * 3), + self._nlist(self.nloc * self.nnei), + self.natoms, + n_a_sel=self.nnei, + n_r_sel=0, + ) + ) + + def test_prod_virial_grad_rejects_rij_width_mismatch(self) -> None: + with self.assertRaisesRegex( + tf.errors.InvalidArgumentError, r"dim of rij should be nnei \* 3" + ): + self.sess.run( + op_grads_module.prod_virial_se_a_grad( + self._floats(9), + self._floats(self.nloc * self.ndescrpt), + self._floats(self.nloc * self.ndescrpt * 3), + self._floats(self.nloc * self.nnei * 3 - 1), + self._nlist(self.nloc * self.nnei), + self.natoms, + n_a_sel=self.nnei, + n_r_sel=0, + ) + ) + + def test_prod_virial_r_grad_rejects_partial_net_deriv_atom(self) -> None: + with self.assertRaisesRegex( + tf.errors.InvalidArgumentError, + r"net deriv width 3 should be divisible by nloc 2", + ): + self.sess.run( + op_grads_module.prod_virial_se_r_grad( + self._floats(9), + self._floats(self.nloc * self.nnei + 1), + self._floats(self.nloc * self.nnei * 3), + self._floats(self.nloc * self.nnei * 3), + self._nlist(self.nloc * self.nnei), + self.natoms, + ) + ) + + +if __name__ == "__main__": + tf.test.main()