From 788f1f51d81e0ed8f7f49b8b6cd5513a4d3631ac Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Thu, 30 Jul 2026 10:40:43 +0800 Subject: [PATCH] fix(ops): validate tabulation tensor shapes Reject malformed TensorFlow and PyTorch tabulation buffers before native kernels can index outside their logical storage. Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- source/lib/include/tabulate_validation.h | 100 +++++++ source/op/pt/tabulate_multi_device.cc | 116 +++++++ source/op/tf/tabulate_multi_device.cc | 283 ++++++++++++++++++ .../tests/pt/test_tabulate_fusion_se_atten.py | 72 +++++ .../tf/test_tabulate_shape_validation.py | 115 +++++++ 5 files changed, 686 insertions(+) create mode 100644 source/lib/include/tabulate_validation.h create mode 100644 source/tests/tf/test_tabulate_shape_validation.py diff --git a/source/lib/include/tabulate_validation.h b/source/lib/include/tabulate_validation.h new file mode 100644 index 0000000000..ec70c2319a --- /dev/null +++ b/source/lib/include/tabulate_validation.h @@ -0,0 +1,100 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +#pragma once + +#include +#include +#include +#include + +namespace deepmd { + +// Multiply non-negative tensor dimensions without invoking signed overflow. +inline bool tabulate_checked_product(const int64_t lhs, + const int64_t rhs, + int64_t& product) { + if (lhs < 0 || rhs < 0 || + (rhs != 0 && lhs > std::numeric_limits::max() / rhs)) { + return false; + } + product = lhs * rhs; + return true; +} + +// Validate the five table metadata values consumed by the native tabulation +// kernels and reproduce the spline-row count used by the table generator. +// Keeping this calculation shared prevents the TensorFlow and PyTorch wrappers +// from accepting different raw-buffer contracts. +template +bool tabulate_required_table_rows(const FPTYPE* table_info, + const bool symmetric_range, + int64_t& required_rows, + std::string& error) { + const double lower = static_cast(table_info[0]); + const double upper = static_cast(table_info[1]); + const double max = static_cast(table_info[2]); + const double stride0 = static_cast(table_info[3]); + const double stride1 = static_cast(table_info[4]); + if (!std::isfinite(lower) || !std::isfinite(upper) || !std::isfinite(max) || + !std::isfinite(stride0) || !std::isfinite(stride1)) { + error = "table_info values must be finite"; + return false; + } + if (stride0 <= 0.0 || stride1 <= 0.0) { + error = "table_info strides must be positive"; + return false; + } + + const double min = symmetric_range ? -max : lower; + if (min > lower || lower > upper || upper > max) { + error = symmetric_range + ? "table_info must satisfy -max <= lower <= upper <= max" + : "table_info must satisfy lower <= upper <= max"; + return false; + } + + const double lower_tail = symmetric_range ? (lower - min) / stride1 : 0.0; + const double middle = (upper - lower) / stride0; + const double upper_tail = (max - upper) / stride1; + const double total_intervals = lower_tail + middle + upper_tail; + const double max_segment = + static_cast(std::numeric_limits::max()); + if (!std::isfinite(lower_tail) || !std::isfinite(middle) || + !std::isfinite(upper_tail) || !std::isfinite(total_intervals) || + total_intervals > max_segment) { + error = "table_info describes too many spline intervals"; + return false; + } + + // The Python table builder converts the sum to an integer once, which is + // observably different from truncating each range separately for SE-T. + required_rows = static_cast(total_intervals); + if (required_rows <= 0) { + error = "table_info must describe at least one spline interval"; + return false; + } + return true; +} + +// Convert the validated row count into the flattened coefficient count while +// guarding the multiplication used by both framework wrappers. +inline bool tabulate_required_table_elements(const int64_t required_rows, + const int64_t last_layer_size, + int64_t& required_elements, + std::string& error) { + constexpr int64_t coefficients_per_feature = 6; + if (required_rows <= 0 || last_layer_size <= 0) { + error = "table dimensions must be positive"; + return false; + } + int64_t feature_elements = 0; + if (!tabulate_checked_product(last_layer_size, coefficients_per_feature, + feature_elements) || + !tabulate_checked_product(required_rows, feature_elements, + required_elements)) { + error = "required table size exceeds the supported integer range"; + return false; + } + return true; +} + +} // namespace deepmd diff --git a/source/op/pt/tabulate_multi_device.cc b/source/op/pt/tabulate_multi_device.cc index cede1d03d9..183240a575 100644 --- a/source/op/pt/tabulate_multi_device.cc +++ b/source/op/pt/tabulate_multi_device.cc @@ -1,10 +1,12 @@ // SPDX-License-Identifier: LGPL-3.0-or-later #include +#include #include #include #include "tabulate.h" +#include "tabulate_validation.h" #if defined(GOOGLE_CUDA) || defined(TENSORFLOW_USE_ROCM) #include "device.h" @@ -18,6 +20,109 @@ void GetTensorDevice(const torch::Tensor& t, std::string& str) { } } +void CheckTabulateDataTensor(const torch::Tensor& tensor, + const torch::Tensor& table_tensor, + const char* name) { + TORCH_CHECK(tensor.scalar_type() == table_tensor.scalar_type(), name, + " must have the same dtype as table"); + TORCH_CHECK(tensor.device() == table_tensor.device(), name, + " must be on the same device as table"); + TORCH_CHECK(tensor.is_contiguous(), name, " must be contiguous"); +} + +template +void CheckTabulateTable(const torch::Tensor& table_tensor, + const torch::Tensor& table_info_tensor, + const int64_t last_layer_size, + const bool symmetric_range) { + TORCH_CHECK(table_tensor.dim() == 2, "table must be rank 2"); + TORCH_CHECK(table_tensor.scalar_type() == torch::kFloat || + table_tensor.scalar_type() == torch::kDouble, + "table must use float32 or float64"); + TORCH_CHECK(table_tensor.device().is_cpu() || table_tensor.device().is_cuda(), + "table must be on a CPU or CUDA/ROCm device"); + TORCH_CHECK(table_tensor.is_contiguous(), "table must be contiguous"); + TORCH_CHECK(last_layer_size > 0, "last_layer_size must be positive"); + TORCH_CHECK(table_info_tensor.device().is_cpu(), + "table_info must be on the CPU"); + TORCH_CHECK(table_info_tensor.scalar_type() == table_tensor.scalar_type(), + "table_info must have the same dtype as table"); + TORCH_CHECK(table_info_tensor.is_contiguous(), + "table_info must be contiguous"); + TORCH_CHECK(table_info_tensor.numel() >= 5, + "table_info must contain at least 5 values"); + + int64_t required_rows = 0; + std::string error; + TORCH_CHECK(deepmd::tabulate_required_table_rows( + table_info_tensor.data_ptr(), symmetric_range, + required_rows, error), + error); + int64_t required_elements = 0; + TORCH_CHECK(deepmd::tabulate_required_table_elements( + required_rows, last_layer_size, required_elements, error), + error); + TORCH_CHECK(table_tensor.numel() >= required_elements, + "table does not contain enough coefficients for table_info and " + "last_layer_size"); +} + +template +void CheckTabulateSeAInputs(const torch::Tensor& table_tensor, + const torch::Tensor& table_info_tensor, + const torch::Tensor& em_x_tensor, + const torch::Tensor& em_tensor, + const torch::Tensor& two_embed_tensor, + const int64_t last_layer_size) { + CheckTabulateTable(table_tensor, table_info_tensor, last_layer_size, + false); + TORCH_CHECK(em_tensor.dim() == 3 && em_tensor.size(2) == 4, + "em must have shape [nloc, nnei, 4]"); + const int64_t neighbor_count = em_tensor.numel() / 4; + TORCH_CHECK(em_x_tensor.dim() == 2 && em_x_tensor.numel() == neighbor_count, + "em_x must be rank 2 and contain nloc * nnei values"); + CheckTabulateDataTensor(em_x_tensor, table_tensor, "em_x"); + CheckTabulateDataTensor(em_tensor, table_tensor, "em"); + if (two_embed_tensor.defined()) { + TORCH_CHECK(two_embed_tensor.dim() == 2, "two_embed must be rank 2"); + int64_t expected_two_embed_elements = 0; + TORCH_CHECK( + deepmd::tabulate_checked_product(neighbor_count, last_layer_size, + expected_two_embed_elements), + "two_embed element count exceeds the supported integer range"); + TORCH_CHECK(two_embed_tensor.numel() == expected_two_embed_elements, + "two_embed must contain nloc * nnei * last_layer_size values"); + CheckTabulateDataTensor(two_embed_tensor, table_tensor, "two_embed"); + } +} + +template +void CheckTabulateSeTInputs(const torch::Tensor& table_tensor, + const torch::Tensor& table_info_tensor, + const torch::Tensor& em_x_tensor, + const torch::Tensor& em_tensor, + const int64_t last_layer_size) { + CheckTabulateTable(table_tensor, table_info_tensor, last_layer_size, + true); + TORCH_CHECK(em_tensor.dim() == 3, "em must be rank 3"); + TORCH_CHECK( + em_x_tensor.dim() == 2 && em_x_tensor.numel() == em_tensor.numel(), + "em_x must be rank 2 and contain the same number of values as em"); + CheckTabulateDataTensor(em_x_tensor, table_tensor, "em_x"); + CheckTabulateDataTensor(em_tensor, table_tensor, "em"); +} + +template +void CheckTabulateSeRInputs(const torch::Tensor& table_tensor, + const torch::Tensor& table_info_tensor, + const torch::Tensor& em_tensor, + const int64_t last_layer_size) { + CheckTabulateTable(table_tensor, table_info_tensor, last_layer_size, + false); + TORCH_CHECK(em_tensor.dim() == 2, "em must be rank 2"); + CheckTabulateDataTensor(em_tensor, table_tensor, "em"); +} + template void TabulateFusionSeAForward(const torch::Tensor& table_tensor, const torch::Tensor& table_info_tensor, @@ -776,6 +881,8 @@ class TabulateFusionSeAOp const torch::Tensor& em_x_tensor, const torch::Tensor& em_tensor, int64_t last_layer_size) { + CheckTabulateSeAInputs(table_tensor, table_info_tensor, em_x_tensor, + em_tensor, at::Tensor(), last_layer_size); // allocate output tensors auto options = torch::TensorOptions() .dtype(table_tensor.dtype()) @@ -962,6 +1069,9 @@ class TabulateFusionSeAttenOp const torch::Tensor& two_embed_tensor, int64_t last_layer_size, bool is_sorted) { + CheckTabulateSeAInputs(table_tensor, table_info_tensor, em_x_tensor, + em_tensor, two_embed_tensor, + last_layer_size); // allocate output tensors auto options = torch::TensorOptions() .dtype(table_tensor.dtype()) @@ -1130,6 +1240,8 @@ class TabulateFusionSeTOp const torch::Tensor& em_x_tensor, const torch::Tensor& em_tensor, int64_t last_layer_size) { + CheckTabulateSeTInputs(table_tensor, table_info_tensor, em_x_tensor, + em_tensor, last_layer_size); // allocate output tensors auto options = torch::TensorOptions() .dtype(table_tensor.dtype()) @@ -1283,6 +1395,8 @@ class TabulateFusionSeROp const torch::Tensor& table_info_tensor, const torch::Tensor& em_tensor, int64_t last_layer_size) { + CheckTabulateSeRInputs(table_tensor, table_info_tensor, em_tensor, + last_layer_size); // allocate output tensors auto options = torch::TensorOptions() .dtype(table_tensor.dtype()) @@ -1441,6 +1555,8 @@ class TabulateFusionSeTTebdOp const torch::Tensor& em_x_tensor, const torch::Tensor& em_tensor, int64_t last_layer_size) { + CheckTabulateSeTInputs(table_tensor, table_info_tensor, em_x_tensor, + em_tensor, last_layer_size); // allocate output tensors auto options = torch::TensorOptions() .dtype(table_tensor.dtype()) diff --git a/source/op/tf/tabulate_multi_device.cc b/source/op/tf/tabulate_multi_device.cc index 174faf5213..134a7b8ec8 100644 --- a/source/op/tf/tabulate_multi_device.cc +++ b/source/op/tf/tabulate_multi_device.cc @@ -2,6 +2,7 @@ // SPDX-License-Identifier: LGPL-3.0-or-later #include "custom_op.h" #include "tabulate.h" +#include "tabulate_validation.h" REGISTER_OP("TabulateFusion") .Attr("T: {float, double} = DT_DOUBLE") @@ -162,6 +163,130 @@ REGISTER_OP("TabulateFusionSeRGradGrad") .Input("descriptor: T") .Output("dz_dy: T"); +template +deepmd::tf_compat::Status ValidateTabulateTable(const Tensor& table_tensor, + const Tensor& table_info_tensor, + const int64_t last_layer_size, + const bool symmetric_range) { + if (table_tensor.dims() != 2) { + return deepmd::tf_compat::InvalidArgument("table must be rank 2"); + } + if (last_layer_size <= 0) { + return deepmd::tf_compat::InvalidArgument( + "last_layer_size must be positive"); + } + if (table_info_tensor.NumElements() < 5) { + return deepmd::tf_compat::InvalidArgument( + "table_info must contain at least 5 values"); + } + int64_t required_rows = 0; + std::string error; + if (!deepmd::tabulate_required_table_rows( + table_info_tensor.flat().data(), symmetric_range, + required_rows, error)) { + return deepmd::tf_compat::InvalidArgument(error); + } + int64_t required_elements = 0; + if (!deepmd::tabulate_required_table_elements(required_rows, last_layer_size, + required_elements, error)) { + return deepmd::tf_compat::InvalidArgument(error); + } + if (table_tensor.NumElements() < required_elements) { + return deepmd::tf_compat::InvalidArgument( + "table does not contain enough coefficients for table_info and " + "last_layer_size"); + } + return deepmd::tf_compat::Status(); +} + +template +deepmd::tf_compat::Status ValidateTabulateSeAInputs( + const Tensor& table_tensor, + const Tensor& table_info_tensor, + const Tensor& em_x_tensor, + const Tensor& em_tensor, + const Tensor* two_embed_tensor, + const int64_t last_layer_size) { + auto status = ValidateTabulateTable(table_tensor, table_info_tensor, + last_layer_size, false); + if (!status.ok()) { + return status; + } + if (em_tensor.dims() != 3 || em_tensor.dim_size(2) != 4) { + return deepmd::tf_compat::InvalidArgument( + "em must have shape [nloc, nnei, 4]"); + } + const int64_t neighbor_count = em_tensor.NumElements() / 4; + if (em_x_tensor.dims() != 2 || em_x_tensor.NumElements() != neighbor_count) { + return deepmd::tf_compat::InvalidArgument( + "em_x must be rank 2 and contain nloc * nnei values"); + } + if (two_embed_tensor != nullptr) { + int64_t expected_two_embed_elements = 0; + if (!deepmd::tabulate_checked_product(neighbor_count, last_layer_size, + expected_two_embed_elements)) { + return deepmd::tf_compat::InvalidArgument( + "two_embed element count exceeds the supported integer range"); + } + if (two_embed_tensor->dims() != 2 || + two_embed_tensor->NumElements() != expected_two_embed_elements) { + return deepmd::tf_compat::InvalidArgument( + "two_embed must be rank 2 and contain nloc * nnei * " + "last_layer_size values"); + } + } + return deepmd::tf_compat::Status(); +} + +template +deepmd::tf_compat::Status ValidateTabulateSeTInputs( + const Tensor& table_tensor, + const Tensor& table_info_tensor, + const Tensor& em_x_tensor, + const Tensor& em_tensor, + const int64_t last_layer_size) { + auto status = ValidateTabulateTable(table_tensor, table_info_tensor, + last_layer_size, true); + if (!status.ok()) { + return status; + } + if (em_tensor.dims() != 3) { + return deepmd::tf_compat::InvalidArgument("em must be rank 3"); + } + if (em_x_tensor.dims() != 2 || + em_x_tensor.NumElements() != em_tensor.NumElements()) { + return deepmd::tf_compat::InvalidArgument( + "em_x must be rank 2 and contain the same number of values as em"); + } + return deepmd::tf_compat::Status(); +} + +template +deepmd::tf_compat::Status ValidateTabulateSeRInputs( + const Tensor& table_tensor, + const Tensor& table_info_tensor, + const Tensor& em_tensor, + const int64_t last_layer_size) { + auto status = ValidateTabulateTable(table_tensor, table_info_tensor, + last_layer_size, false); + if (!status.ok()) { + return status; + } + if (em_tensor.dims() != 2) { + return deepmd::tf_compat::InvalidArgument("em must be rank 2"); + } + return deepmd::tf_compat::Status(); +} + +deepmd::tf_compat::Status ValidateTensorShape(const Tensor& tensor, + const TensorShape& expected, + const char* name) { + if (tensor.shape() != expected) { + return deepmd::tf_compat::InvalidArgument(name, " has an unexpected shape"); + } + return deepmd::tf_compat::Status(); +} + template class TabulateFusionSeAOp : public OpKernel { public: @@ -182,6 +307,9 @@ class TabulateFusionSeAOp : public OpKernel { const Tensor& table_info_tensor = context->input(context_input_index++); const Tensor& em_x_tensor = context->input(context_input_index++); const Tensor& em_tensor = context->input(context_input_index++); + OP_REQUIRES_OK(context, ValidateTabulateSeAInputs( + table_tensor, table_info_tensor, em_x_tensor, + em_tensor, nullptr, last_layer_size)); // set size of the sample OP_REQUIRES(context, (table_tensor.shape().dims() == 2), deepmd::tf_compat::InvalidArgument("Dim of table should be 2")); @@ -245,6 +373,22 @@ class TabulateFusionSeAGradOp : public OpKernel { const Tensor& dy_tensor = context->input(context_input_index++); const Tensor& descriptor_tensor = context->input(context_input_index++); + OP_REQUIRES(context, descriptor_tensor.dims() == 3, + deepmd::tf_compat::InvalidArgument( + "descriptor must have shape [nloc, 4, last_layer_size]")); + const int64_t validated_last_layer_size = descriptor_tensor.dim_size(2); + OP_REQUIRES_OK(context, ValidateTabulateSeAInputs( + table_tensor, table_info_tensor, em_x_tensor, + em_tensor, nullptr, validated_last_layer_size)); + TensorShape expected_descriptor_shape; + expected_descriptor_shape.AddDim(em_tensor.dim_size(0)); + expected_descriptor_shape.AddDim(4); + expected_descriptor_shape.AddDim(validated_last_layer_size); + OP_REQUIRES_OK( + context, ValidateTensorShape(descriptor_tensor, + expected_descriptor_shape, "descriptor")); + OP_REQUIRES_OK(context, ValidateTensorShape( + dy_tensor, expected_descriptor_shape, "dy")); // set size of the sample OP_REQUIRES(context, (dy_tensor.shape().dims() == 3), deepmd::tf_compat::InvalidArgument("Dim of table should be 3")); @@ -309,6 +453,26 @@ class TabulateFusionSeAGradGradOp : public OpKernel { const Tensor& dz_dy_dem_x_tensor = context->input(context_input_index++); const Tensor& dz_dy_dem_tensor = context->input(context_input_index++); const Tensor& descriptor_tensor = context->input(context_input_index++); + OP_REQUIRES(context, descriptor_tensor.dims() == 3, + deepmd::tf_compat::InvalidArgument( + "descriptor must have shape [nloc, 4, last_layer_size]")); + const int64_t validated_last_layer_size = descriptor_tensor.dim_size(2); + OP_REQUIRES_OK(context, ValidateTabulateSeAInputs( + table_tensor, table_info_tensor, em_x_tensor, + em_tensor, nullptr, validated_last_layer_size)); + TensorShape expected_descriptor_shape; + expected_descriptor_shape.AddDim(em_tensor.dim_size(0)); + expected_descriptor_shape.AddDim(4); + expected_descriptor_shape.AddDim(validated_last_layer_size); + OP_REQUIRES_OK( + context, ValidateTensorShape(descriptor_tensor, + expected_descriptor_shape, "descriptor")); + OP_REQUIRES_OK(context, + ValidateTensorShape(dz_dy_dem_x_tensor, em_x_tensor.shape(), + "dz_dy_dem_x")); + OP_REQUIRES_OK( + context, + ValidateTensorShape(dz_dy_dem_tensor, em_tensor.shape(), "dz_dy_dem")); // set size of the sample OP_REQUIRES(context, (dz_dy_dem_x_tensor.shape().dims() == 2), deepmd::tf_compat::InvalidArgument("Dim of input should be 2")); @@ -379,6 +543,9 @@ class TabulateFusionSeAttenOp : public OpKernel { const Tensor& em_x_tensor = context->input(context_input_index++); const Tensor& em_tensor = context->input(context_input_index++); const Tensor& two_embed_tensor = context->input(context_input_index++); + OP_REQUIRES_OK(context, ValidateTabulateSeAInputs( + table_tensor, table_info_tensor, em_x_tensor, + em_tensor, &two_embed_tensor, last_layer_size)); // set size of the sample OP_REQUIRES(context, (table_tensor.shape().dims() == 2), deepmd::tf_compat::InvalidArgument("Dim of table should be 2")); @@ -450,6 +617,23 @@ class TabulateFusionSeAttenGradOp : public OpKernel { const Tensor& dy_tensor = context->input(context_input_index++); const Tensor& descriptor_tensor = context->input(context_input_index++); + OP_REQUIRES(context, descriptor_tensor.dims() == 3, + deepmd::tf_compat::InvalidArgument( + "descriptor must have shape [nloc, 4, last_layer_size]")); + const int64_t validated_last_layer_size = descriptor_tensor.dim_size(2); + OP_REQUIRES_OK(context, + ValidateTabulateSeAInputs( + table_tensor, table_info_tensor, em_x_tensor, em_tensor, + &two_embed_tensor, validated_last_layer_size)); + TensorShape expected_descriptor_shape; + expected_descriptor_shape.AddDim(em_tensor.dim_size(0)); + expected_descriptor_shape.AddDim(4); + expected_descriptor_shape.AddDim(validated_last_layer_size); + OP_REQUIRES_OK( + context, ValidateTensorShape(descriptor_tensor, + expected_descriptor_shape, "descriptor")); + OP_REQUIRES_OK(context, ValidateTensorShape( + dy_tensor, expected_descriptor_shape, "dy")); // set size of the sample OP_REQUIRES(context, (dy_tensor.shape().dims() == 3), deepmd::tf_compat::InvalidArgument("Dim of table should be 3")); @@ -526,6 +710,30 @@ class TabulateFusionSeAttenGradGradOp : public OpKernel { const Tensor& dz_dy_dem_tensor = context->input(context_input_index++); const Tensor& dz_dy_dtwo_tensor = context->input(context_input_index++); const Tensor& descriptor_tensor = context->input(context_input_index++); + OP_REQUIRES(context, descriptor_tensor.dims() == 3, + deepmd::tf_compat::InvalidArgument( + "descriptor must have shape [nloc, 4, last_layer_size]")); + const int64_t validated_last_layer_size = descriptor_tensor.dim_size(2); + OP_REQUIRES_OK(context, + ValidateTabulateSeAInputs( + table_tensor, table_info_tensor, em_x_tensor, em_tensor, + &two_embed_tensor, validated_last_layer_size)); + TensorShape expected_descriptor_shape; + expected_descriptor_shape.AddDim(em_tensor.dim_size(0)); + expected_descriptor_shape.AddDim(4); + expected_descriptor_shape.AddDim(validated_last_layer_size); + OP_REQUIRES_OK( + context, ValidateTensorShape(descriptor_tensor, + expected_descriptor_shape, "descriptor")); + OP_REQUIRES_OK(context, + ValidateTensorShape(dz_dy_dem_x_tensor, em_x_tensor.shape(), + "dz_dy_dem_x")); + OP_REQUIRES_OK( + context, + ValidateTensorShape(dz_dy_dem_tensor, em_tensor.shape(), "dz_dy_dem")); + OP_REQUIRES_OK(context, + ValidateTensorShape(dz_dy_dtwo_tensor, + two_embed_tensor.shape(), "dz_dy_dtwo")); // set size of the sample OP_REQUIRES(context, (dz_dy_dem_x_tensor.shape().dims() == 2), deepmd::tf_compat::InvalidArgument("Dim of input should be 2")); @@ -594,6 +802,9 @@ class TabulateFusionSeTOp : public OpKernel { const Tensor& table_info_tensor = context->input(context_input_index++); const Tensor& em_x_tensor = context->input(context_input_index++); const Tensor& em_tensor = context->input(context_input_index++); + OP_REQUIRES_OK(context, ValidateTabulateSeTInputs( + table_tensor, table_info_tensor, em_x_tensor, + em_tensor, last_layer_size)); // set size of the sample OP_REQUIRES(context, (table_tensor.shape().dims() == 2), deepmd::tf_compat::InvalidArgument("Dim of table should be 2")); @@ -657,6 +868,21 @@ class TabulateFusionSeTGradOp : public OpKernel { const Tensor& em_tensor = context->input(context_input_index++); const Tensor& dy_tensor = context->input(context_input_index++); const Tensor& descriptor_tensor = context->input(context_input_index++); + OP_REQUIRES(context, descriptor_tensor.dims() == 2, + deepmd::tf_compat::InvalidArgument( + "descriptor must have shape [nloc, last_layer_size]")); + const int64_t validated_last_layer_size = descriptor_tensor.dim_size(1); + OP_REQUIRES_OK(context, ValidateTabulateSeTInputs( + table_tensor, table_info_tensor, em_x_tensor, + em_tensor, validated_last_layer_size)); + TensorShape expected_descriptor_shape; + expected_descriptor_shape.AddDim(em_tensor.dim_size(0)); + expected_descriptor_shape.AddDim(validated_last_layer_size); + OP_REQUIRES_OK( + context, ValidateTensorShape(descriptor_tensor, + expected_descriptor_shape, "descriptor")); + OP_REQUIRES_OK(context, ValidateTensorShape( + dy_tensor, expected_descriptor_shape, "dy")); // set size of the sample OP_REQUIRES( context, (dy_tensor.shape().dims() == 2), @@ -718,6 +944,25 @@ class TabulateFusionSeTGradGradOp : public OpKernel { const Tensor& dz_dy_dem_x_tensor = context->input(context_input_index++); const Tensor& dz_dy_dem_tensor = context->input(context_input_index++); const Tensor& descriptor_tensor = context->input(context_input_index++); + OP_REQUIRES(context, descriptor_tensor.dims() == 2, + deepmd::tf_compat::InvalidArgument( + "descriptor must have shape [nloc, last_layer_size]")); + const int64_t validated_last_layer_size = descriptor_tensor.dim_size(1); + OP_REQUIRES_OK(context, ValidateTabulateSeTInputs( + table_tensor, table_info_tensor, em_x_tensor, + em_tensor, validated_last_layer_size)); + TensorShape expected_descriptor_shape; + expected_descriptor_shape.AddDim(em_tensor.dim_size(0)); + expected_descriptor_shape.AddDim(validated_last_layer_size); + OP_REQUIRES_OK( + context, ValidateTensorShape(descriptor_tensor, + expected_descriptor_shape, "descriptor")); + OP_REQUIRES_OK(context, + ValidateTensorShape(dz_dy_dem_x_tensor, em_x_tensor.shape(), + "dz_dy_dem_x")); + OP_REQUIRES_OK( + context, + ValidateTensorShape(dz_dy_dem_tensor, em_tensor.shape(), "dz_dy_dem")); // set size of the sample OP_REQUIRES(context, (dz_dy_dem_x_tensor.shape().dims() == 2), deepmd::tf_compat::InvalidArgument("Dim of input should be 2")); @@ -782,6 +1027,9 @@ class TabulateFusionSeROp : public OpKernel { const Tensor& table_tensor = context->input(context_input_index++); const Tensor& table_info_tensor = context->input(context_input_index++); const Tensor& em_tensor = context->input(context_input_index++); + OP_REQUIRES_OK(context, ValidateTabulateSeRInputs( + table_tensor, table_info_tensor, em_tensor, + last_layer_size)); // set size of the sample OP_REQUIRES(context, (table_tensor.shape().dims() == 2), deepmd::tf_compat::InvalidArgument("Dim of table should be 2")); @@ -839,6 +1087,23 @@ class TabulateFusionSeRGradOp : public OpKernel { const Tensor& em_tensor = context->input(context_input_index++); const Tensor& dy_tensor = context->input(context_input_index++); const Tensor& descriptor_tensor = context->input(context_input_index++); + OP_REQUIRES(context, descriptor_tensor.dims() == 3, + deepmd::tf_compat::InvalidArgument( + "descriptor must have shape [nloc, nnei, " + "last_layer_size]")); + const int64_t validated_last_layer_size = descriptor_tensor.dim_size(2); + OP_REQUIRES_OK(context, ValidateTabulateSeRInputs( + table_tensor, table_info_tensor, em_tensor, + validated_last_layer_size)); + TensorShape expected_descriptor_shape; + expected_descriptor_shape.AddDim(em_tensor.dim_size(0)); + expected_descriptor_shape.AddDim(em_tensor.dim_size(1)); + expected_descriptor_shape.AddDim(validated_last_layer_size); + OP_REQUIRES_OK( + context, ValidateTensorShape(descriptor_tensor, + expected_descriptor_shape, "descriptor")); + OP_REQUIRES_OK(context, ValidateTensorShape( + dy_tensor, expected_descriptor_shape, "dy")); // set size of the sample OP_REQUIRES(context, (dy_tensor.shape().dims() == 3), deepmd::tf_compat::InvalidArgument("Dim of table should be 3")); @@ -888,6 +1153,24 @@ class TabulateFusionSeRGradGradOp : public OpKernel { const Tensor& em_tensor = context->input(context_input_index++); const Tensor& dz_dy_dem_tensor = context->input(context_input_index++); const Tensor& descriptor_tensor = context->input(context_input_index++); + OP_REQUIRES(context, descriptor_tensor.dims() == 3, + deepmd::tf_compat::InvalidArgument( + "descriptor must have shape [nloc, nnei, " + "last_layer_size]")); + const int64_t validated_last_layer_size = descriptor_tensor.dim_size(2); + OP_REQUIRES_OK(context, ValidateTabulateSeRInputs( + table_tensor, table_info_tensor, em_tensor, + validated_last_layer_size)); + TensorShape expected_descriptor_shape; + expected_descriptor_shape.AddDim(em_tensor.dim_size(0)); + expected_descriptor_shape.AddDim(em_tensor.dim_size(1)); + expected_descriptor_shape.AddDim(validated_last_layer_size); + OP_REQUIRES_OK( + context, ValidateTensorShape(descriptor_tensor, + expected_descriptor_shape, "descriptor")); + OP_REQUIRES_OK( + context, + ValidateTensorShape(dz_dy_dem_tensor, em_tensor.shape(), "dz_dy_dem")); // set size of the sample OP_REQUIRES(context, (dz_dy_dem_tensor.shape().dims() == 2), deepmd::tf_compat::InvalidArgument("Dim of input should be 2")); diff --git a/source/tests/pt/test_tabulate_fusion_se_atten.py b/source/tests/pt/test_tabulate_fusion_se_atten.py index f46d0b1761..e6f851552b 100644 --- a/source/tests/pt/test_tabulate_fusion_se_atten.py +++ b/source/tests/pt/test_tabulate_fusion_se_atten.py @@ -1644,6 +1644,78 @@ def test_second_order_backward(self) -> None: (self.em_x_tensor, self.em_tensor, self.two_embed_tensor), ) + def test_rejects_mismatched_native_buffer_shapes(self) -> None: + invalid_inputs = ( + ( + self.em_x_tensor[:, :-1].contiguous(), + self.em_tensor, + self.two_embed_tensor, + "em_x must be rank 2", + ), + ( + self.em_x_tensor, + self.em_tensor[..., :3].contiguous(), + self.two_embed_tensor, + "em must have shape", + ), + ( + self.em_x_tensor, + self.em_tensor, + self.two_embed_tensor.reshape(-1)[:-1].reshape(1, -1), + "two_embed must contain", + ), + ) + for em_x, em, two_embed, message in invalid_inputs: + with ( + self.subTest(message=message), + self.assertRaisesRegex(RuntimeError, message), + ): + torch.ops.deepmd.tabulate_fusion_se_atten( + self.table_tensor, + self.table_info_tensor, + em_x, + em, + two_embed, + self.last_layer_size, + self.is_sorted, + ) + + def test_accepts_flattened_em_x_layout(self) -> None: + result = torch.ops.deepmd.tabulate_fusion_se_atten( + self.table_tensor, + self.table_info_tensor, + self.em_x_tensor.reshape(-1, 1), + self.em_tensor, + self.two_embed_tensor, + self.last_layer_size, + self.is_sorted, + ) + self.assertEqual(result[0].shape, self.expected_descriptor_tensor.shape) + + def test_rejects_short_table_buffers(self) -> None: + with self.assertRaisesRegex(RuntimeError, "table_info must contain"): + torch.ops.deepmd.tabulate_fusion_se_atten( + self.table_tensor, + self.table_info_tensor[:4], + self.em_x_tensor, + self.em_tensor, + self.two_embed_tensor, + self.last_layer_size, + self.is_sorted, + ) + + short_table = self.table_tensor.reshape(-1)[:-1].reshape(1, -1) + with self.assertRaisesRegex(RuntimeError, "table does not contain enough"): + torch.ops.deepmd.tabulate_fusion_se_atten( + short_table, + self.table_info_tensor, + self.em_x_tensor, + self.em_tensor, + self.two_embed_tensor, + self.last_layer_size, + self.is_sorted, + ) + if __name__ == "__main__": unittest.main() diff --git a/source/tests/tf/test_tabulate_shape_validation.py b/source/tests/tf/test_tabulate_shape_validation.py new file mode 100644 index 0000000000..c090f2a492 --- /dev/null +++ b/source/tests/tf/test_tabulate_shape_validation.py @@ -0,0 +1,115 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import unittest + +import numpy as np + +from deepmd.tf.env import ( + op_module, + tf, +) + + +class TestTabulateShapeValidation(unittest.TestCase): + def setUp(self) -> None: + self.table = tf.constant(np.zeros((2, 12)), dtype=tf.float64) + self.table_info = tf.constant([0, 1, 2, 1, 1, -1], dtype=tf.float64) + self.em_x = tf.constant([[0.25, 0.5, 0.75]], dtype=tf.float64) + self.em = tf.constant(np.zeros((1, 3, 4)), dtype=tf.float64) + + def test_rejects_short_two_embed(self) -> None: + descriptor = op_module.tabulate_fusion_se_atten( + self.table, + self.table_info, + self.em_x, + self.em, + tf.constant([[1.0]], dtype=tf.float64), + last_layer_size=2, + is_sorted=True, + ) + with ( + self.assertRaisesRegex( + tf.errors.InvalidArgumentError, "two_embed must be rank 2" + ), + tf.Session() as sess, + ): + sess.run(descriptor) + + def test_accepts_flattened_em_x_layout(self) -> None: + descriptor = op_module.tabulate_fusion_se_a( + self.table, + self.table_info, + tf.reshape(self.em_x, [-1, 1]), + self.em, + last_layer_size=2, + ) + with tf.Session() as sess: + self.assertEqual(sess.run(descriptor).shape, (1, 4, 2)) + + def test_rejects_mismatched_em_x(self) -> None: + descriptor = op_module.tabulate_fusion_se_a( + self.table, + self.table_info, + self.em_x[:, :2], + self.em, + last_layer_size=2, + ) + with ( + self.assertRaisesRegex( + tf.errors.InvalidArgumentError, "em_x must be rank 2" + ), + tf.Session() as sess, + ): + sess.run(descriptor) + + def test_rejects_mismatched_gradient_shape(self) -> None: + gradients = op_module.tabulate_fusion_se_a_grad( + self.table, + self.table_info, + self.em_x, + self.em, + tf.zeros((1, 4, 1), dtype=tf.float64), + tf.zeros((1, 4, 2), dtype=tf.float64), + ) + with ( + self.assertRaisesRegex( + tf.errors.InvalidArgumentError, "dy has an unexpected shape" + ), + tf.Session() as sess, + ): + sess.run(gradients) + + def test_rejects_short_table_info(self) -> None: + descriptor = op_module.tabulate_fusion_se_a( + self.table, + self.table_info[:4], + self.em_x, + self.em, + last_layer_size=2, + ) + with ( + self.assertRaisesRegex( + tf.errors.InvalidArgumentError, "table_info must contain" + ), + tf.Session() as sess, + ): + sess.run(descriptor) + + def test_rejects_short_table(self) -> None: + descriptor = op_module.tabulate_fusion_se_a( + self.table[:, :-1], + self.table_info, + self.em_x, + self.em, + last_layer_size=2, + ) + with ( + self.assertRaisesRegex( + tf.errors.InvalidArgumentError, "table does not contain enough" + ), + tf.Session() as sess, + ): + sess.run(descriptor) + + +if __name__ == "__main__": + unittest.main()