From ffd8a5cc6dd5aae9e7921684ff626983b130c156 Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Fri, 17 Jul 2026 06:12:33 +0800 Subject: [PATCH 1/2] fix(jax): materialize default fparam in cxx api Load and validate default frame parameters from JAX SavedModels, materialize stored or caller-provided values for direct and neighbor-list inference, and broadcast one-frame values across multiple frames. Validate TensorFlow C API tensor shapes and sizes before allocation, cast rank explicitly, and report allocation failure before accessing tensor storage. Cover default and override behavior, float and double paths, multiframe broadcasting, neighbor lists, and invalid sizes. Coding-Agent: Codex Codex-Version: codex-cli 0.144.4 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- source/api_cc/include/DeepPotJAX.h | 4 + source/api_cc/src/DeepPotJAX.cc | 112 +++++++++++++++++- source/api_cc/tests/test_deeppot_universal.cc | 88 +++++++++++++- source/tests/infer/gen_fparam_aparam.py | 26 +++- 4 files changed, 221 insertions(+), 9 deletions(-) diff --git a/source/api_cc/include/DeepPotJAX.h b/source/api_cc/include/DeepPotJAX.h index 93dddd8c8f..2fe365389d 100644 --- a/source/api_cc/include/DeepPotJAX.h +++ b/source/api_cc/include/DeepPotJAX.h @@ -269,6 +269,10 @@ class DeepPotJAX : public DeepPotBackend { bool do_message_passing; // has default fparam bool has_default_fparam_; + // Frame parameters embedded in the exported model. The JAX SavedModel + // signatures require a concrete tensor, so the C++ API materializes these + // values when callers omit fparam. + std::vector default_fparam_; // Model-level pair-exclusion keep table (flat (ntypes+1)^2), built once in // init from the exported ``get_pair_exclude_types``. Empty => no exclusion. // The exported ``call_lower_*`` consumes a pre-excluded nlist (decision diff --git a/source/api_cc/src/DeepPotJAX.cc b/source/api_cc/src/DeepPotJAX.cc index 22ccb5246b..54ee07c5f2 100644 --- a/source/api_cc/src/DeepPotJAX.cc +++ b/source/api_cc/src/DeepPotJAX.cc @@ -12,6 +12,7 @@ #include #include #include +#include #include #include #include @@ -439,10 +440,38 @@ inline std::vector get_vector_string( template inline TF_Tensor* create_tensor(const std::vector& data, const std::vector& shape) { + size_t element_count = 1; + for (const int64_t dim : shape) { + if (dim < 0) { + throw deepmd::deepmd_exception( + "Cannot create a tensor with a negative dimension."); + } + const size_t size_dim = static_cast(dim); + if (size_dim != 0 && + element_count > std::numeric_limits::max() / size_dim) { + throw deepmd::deepmd_exception("Tensor element count overflows size_t."); + } + element_count *= size_dim; + } + if (element_count != data.size()) { + throw deepmd::deepmd_exception( + "Tensor shape requires " + std::to_string(element_count) + + " values, but " + std::to_string(data.size()) + " were provided."); + } + if (element_count > std::numeric_limits::max() / sizeof(T)) { + throw deepmd::deepmd_exception("Tensor byte size overflows size_t."); + } + const size_t byte_size = element_count * sizeof(T); TF_Tensor* tensor = - TF_AllocateTensor(get_data_tensor_type(data), shape.data(), shape.size(), - data.size() * sizeof(T)); - memcpy(TF_TensorData(tensor), data.data(), TF_TensorByteSize(tensor)); + TF_AllocateTensor(get_data_tensor_type(data), shape.data(), + static_cast(shape.size()), byte_size); + if (tensor == nullptr) { + throw deepmd::deepmd_exception( + "TensorFlow failed to allocate an input tensor."); + } + if (byte_size != 0) { + memcpy(TF_TensorData(tensor), data.data(), byte_size); + } return tensor; } @@ -509,6 +538,59 @@ inline std::vector make_charge_spin_input( std::to_string(nframes) + " frames)."); } +inline std::vector make_fparam_input( + const std::vector& fparam, + const int dfparam, + const int nframes, + const bool has_default_fparam, + const std::vector& default_fparam) { + if (dfparam == 0) { + if (!fparam.empty()) { + throw deepmd::deepmd_exception( + "fparam was provided, but this model has dim_fparam=0."); + } + return {}; + } + + const std::vector* source = nullptr; + if (!fparam.empty()) { + source = &fparam; + } else if (!default_fparam.empty()) { + source = &default_fparam; + } else if (has_default_fparam) { + throw deepmd::deepmd_exception( + "fparam was omitted, but the model's default_fparam values are not " + "available. Regenerate the JAX SavedModel with an updated version of " + "deepmd-kit or provide fparam explicitly."); + } else { + throw deepmd::deepmd_exception( + "fparam is required for this model but was not provided, and no " + "default_fparam is stored in the model."); + } + + const size_t dim = static_cast(dfparam); + if (static_cast(nframes) > std::numeric_limits::max() / dim) { + throw deepmd::deepmd_exception("fparam element count overflows size_t."); + } + const size_t expected = static_cast(nframes) * dim; + if (source->size() == expected) { + return *source; + } + if (source->size() == dim) { + std::vector result(expected); + for (int ff = 0; ff < nframes; ++ff) { + std::copy(source->begin(), source->end(), + result.begin() + static_cast(ff) * dim); + } + return result; + } + throw deepmd::deepmd_exception( + "fparam has " + std::to_string(source->size()) + + " values but the model expects dim_fparam=" + std::to_string(dfparam) + + " (per frame) or " + std::to_string(expected) + " (for " + + std::to_string(nframes) + " frames)."); +} + template inline void tensor_to_vector(std::vector& result, TFE_TensorHandle* retval, @@ -672,6 +754,26 @@ void deepmd::DeepPotJAX::init(const std::string& model, } catch (tf_function_not_found& e) { has_default_fparam_ = false; } + if (dfparam > 0 && has_default_fparam_) { + try { + default_fparam_ = get_vector(ctx, "get_default_fparam", + func_vector, device, status); + if (static_cast(default_fparam_.size()) != dfparam) { + throw deepmd::deepmd_exception( + "default_fparam length (" + std::to_string(default_fparam_.size()) + + ") does not match dim_fparam (" + std::to_string(dfparam) + ")."); + } + } catch (tf_function_not_found& e) { + default_fparam_.clear(); + std::cerr << "WARNING: Model has has_default_fparam=true but the " + "get_default_fparam function is missing. Empty fparam will " + "not be substituted. Please regenerate the JAX SavedModel " + "with an updated version of deepmd-kit." + << std::endl; + } + } else { + default_fparam_.clear(); + } try { // Model-level pair_exclude_types, exported flat [ti0, tj0, ti1, tj1, ...]. // Fold exclusion into the LAMMPS nlist at ingestion (decision #18/A4); the @@ -749,6 +851,8 @@ void deepmd::DeepPotJAX::compute(std::vector& ener, std::vector coord_double(coord.begin(), coord.end()); std::vector box_double(box.begin(), box.end()); std::vector fparam_double(fparam.begin(), fparam.end()); + fparam_double = make_fparam_input(fparam_double, dfparam, nframes, + has_default_fparam_, default_fparam_); std::vector aparam_double(aparam.begin(), aparam.end()); std::vector charge_spin_double = make_charge_spin_input(charge_spin, dchgspin, nframes, default_chg_spin_); @@ -912,6 +1016,8 @@ void deepmd::DeepPotJAX::compute(std::vector& ener, // float model interface std::vector coord_double(coord.begin(), coord.end()); std::vector fparam_double(fparam.begin(), fparam.end()); + fparam_double = make_fparam_input(fparam_double, dfparam, nframes, + has_default_fparam_, default_fparam_); std::vector aparam_double(aparam.begin(), aparam.end()); std::vector charge_spin_double = make_charge_spin_input(charge_spin, dchgspin, nframes, default_chg_spin_); diff --git a/source/api_cc/tests/test_deeppot_universal.cc b/source/api_cc/tests/test_deeppot_universal.cc index 59aeb647c2..4f147efca5 100644 --- a/source/api_cc/tests/test_deeppot_universal.cc +++ b/source/api_cc/tests/test_deeppot_universal.cc @@ -7,6 +7,7 @@ #include #include #include +#include #include #include "DeepPot.h" @@ -307,6 +308,10 @@ std::vector default_fparam_cases() { {"pytorch_pt2", Backend::PTExpt, "../../tests/infer/fparam_aparam_default.pt2", "../../tests/infer/fparam_aparam_default.expected", 1e-7, 1e-4, + /*supports_float=*/true}, + {"jax_savedmodel", Backend::JAX, + "../../tests/infer/fparam_aparam_default.savedmodel", + "../../tests/infer/fparam_aparam_default.expected", 1e-7, 1e-4, /*supports_float=*/true}}; } @@ -473,6 +478,8 @@ class DefaultFParamDeepPotTest protected: deepmd::DeepPot dp; deepmd_test::DeepPotRef ref; + deepmd_test::DeepPotRef override_ref; + const std::vector override_fparam = {0.5}; void SetUp() override { const auto& param = GetParam(); @@ -484,6 +491,7 @@ class DefaultFParamDeepPotTest ASSERT_TRUE(path_exists(param.ref_path)) << "Reference artifact is not available: " << param.ref_path; ref = load_fparam_ref(param.ref_path); + override_ref = load_expected_ref(param.ref_path, "override"); ref.has_default_fparam = true; dp.init(param.model_path); } @@ -1942,8 +1950,8 @@ TEST_P(DefaultFParamDeepPotTest, ComputeWithEmptyFParamFloat) { } TEST_P(DefaultFParamDeepPotTest, ComputeWithExplicitFParamDouble) { - check_fparam_compute_simple(dp, ref, GetParam().double_tol, - deepmd_test::fparam_value()); + check_fparam_compute_simple(dp, override_ref, GetParam().double_tol, + override_fparam); } TEST_P(DefaultFParamDeepPotTest, ComputeWithExplicitFParamFloat) { @@ -1951,8 +1959,8 @@ TEST_P(DefaultFParamDeepPotTest, ComputeWithExplicitFParamFloat) { GTEST_SKIP() << backend_name(GetParam().backend) << " does not provide float inference coverage."; } - check_fparam_compute_simple(dp, ref, GetParam().float_tol, - deepmd_test::fparam_value()); + check_fparam_compute_simple(dp, override_ref, GetParam().float_tol, + override_fparam); } TEST_P(DefaultFParamDeepPotTest, LmpNlistWithEmptyFParamDouble) { @@ -1969,6 +1977,78 @@ TEST_P(DefaultFParamDeepPotTest, LmpNlistWithEmptyFParamFloat) { 1); } +TEST_P(DefaultFParamDeepPotTest, LmpNlistWithExplicitFParamDouble) { + check_fparam_lmp_nlist(dp, override_ref, GetParam().double_tol, + override_fparam, false, 1.0, 1); +} + +TEST_P(DefaultFParamDeepPotTest, LmpNlistWithExplicitFParamFloat) { + if (!GetParam().supports_float) { + GTEST_SKIP() << backend_name(GetParam().backend) + << " does not provide float inference coverage."; + } + check_fparam_lmp_nlist(dp, override_ref, GetParam().float_tol, + override_fparam, false, 1.0, 1); +} + +TEST_P(DefaultFParamDeepPotTest, JAXBroadcastsFParamAcrossFrames) { + if (GetParam().backend != Backend::JAX) { + GTEST_SKIP() << "This regression targets JAX SavedModel tensor inputs."; + } + const int nframes = 2; + const std::vector coord = + repeat_values(deepmd_test::deeppot_coord(), nframes); + const std::vector atype = deepmd_test::fparam_aparam_atype(); + const std::vector box = + repeat_values(deepmd_test::deeppot_box(), nframes); + const std::vector aparam = + repeat_values(deepmd_test::aparam_value(), nframes); + const std::vector< + std::pair, const deepmd_test::DeepPotRef*>> + fparam_cases = {{{}, &ref}, {override_fparam, &override_ref}}; + for (const auto& [fparam, expected_ref] : fparam_cases) { + SCOPED_TRACE(fparam.empty() ? "stored default" : "explicit override"); + const std::vector expected_virial = + deepmd_test::total_virial(*expected_ref); + std::vector energy, force, virial; + dp.compute(energy, force, virial, coord, atype, box, fparam, aparam); + + ASSERT_EQ(energy.size(), static_cast(nframes)); + ASSERT_EQ(force.size(), expected_ref->force.size() * nframes); + ASSERT_EQ(virial.size(), 9U * nframes); + for (int ff = 0; ff < nframes; ++ff) { + EXPECT_NEAR(energy[ff], deepmd_test::total_energy(*expected_ref), + GetParam().double_tol); + for (size_t ii = 0; ii < expected_ref->force.size(); ++ii) { + EXPECT_NEAR( + force[static_cast(ff) * expected_ref->force.size() + ii], + expected_ref->force[ii], GetParam().double_tol); + } + for (size_t ii = 0; ii < expected_virial.size(); ++ii) { + EXPECT_NEAR(virial[static_cast(ff) * 9 + ii], + expected_virial[ii], GetParam().double_tol); + } + } + } +} + +TEST_P(DefaultFParamDeepPotTest, JAXRejectsInvalidFParamSize) { + if (GetParam().backend != Backend::JAX) { + GTEST_SKIP() << "This regression targets JAX SavedModel tensor inputs."; + } + const std::vector coord = deepmd_test::deeppot_coord(); + const std::vector atype = deepmd_test::fparam_aparam_atype(); + const std::vector box = deepmd_test::deeppot_box(); + const std::vector aparam = deepmd_test::aparam_value(); + const std::vector invalid_fparam = {0.1, 0.2}; + double energy = 0.0; + std::vector force, virial; + + EXPECT_THROW(dp.compute(energy, force, virial, coord, atype, box, + invalid_fparam, aparam), + deepmd::deepmd_exception); +} + INSTANTIATE_TEST_SUITE_P( AvailableBackends, UniversalDeepPotTest, diff --git a/source/tests/infer/gen_fparam_aparam.py b/source/tests/infer/gen_fparam_aparam.py index 2445b8b304..a530642de8 100644 --- a/source/tests/infer/gen_fparam_aparam.py +++ b/source/tests/infer/gen_fparam_aparam.py @@ -1,10 +1,11 @@ #!/usr/bin/env python3 # SPDX-License-Identifier: LGPL-3.0-or-later -"""Generate fparam_aparam .pt2 test models from pre-committed .pth. +"""Generate fparam_aparam test models from a pre-committed .pth. Converts the pre-committed fparam_aparam_default.pth (checked into git) to -.pt2 format. Produces: +portable inference formats. Produces: - fparam_aparam_default.pt2: model WITH default_fparam + - fparam_aparam_default.savedmodel: JAX model WITH default_fparam - fparam_aparam.pt2: same weights WITHOUT default_fparam - fparam_aparam.pth: same weights WITHOUT default_fparam (torch.jit) Also prints reference values for C++ tests. @@ -30,6 +31,9 @@ def main(): import torch + from deepmd.jax.utils.serialization import ( + deserialize_to_file as jax_deserialize_to_file, + ) from deepmd.pt.model.model import ( get_model, ) @@ -72,6 +76,10 @@ def main(): pt2_default_path, copy.deepcopy(data_default), do_atomic_virial=True ) + savedmodel_default_path = os.path.join(base_dir, "fparam_aparam_default.savedmodel") + print(f"Exporting to {savedmodel_default_path} ...") # noqa: T201 + jax_deserialize_to_file(savedmodel_default_path, copy.deepcopy(data_default)) + # ---- 3. Export fparam_aparam.pt2 and .pth (without default_fparam) ---- config_no_default = copy.deepcopy(config) config_no_default["fitting_net"].pop("default_fparam", None) @@ -135,6 +143,7 @@ def main(): atype = [0, 0, 0, 0, 0, 0] box = np.array([13.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0], dtype=np.float64) fparam_val = np.array([0.25852028], dtype=np.float64) + fparam_override = np.array([0.5], dtype=np.float64) aparam_val = np.array([0.25852028] * 6, dtype=np.float64) e, f, v, ae, av = dp.eval( @@ -177,6 +186,14 @@ def main(): aparam=aparam_val, atomic=True, ) + _e_override, f_override, _v_override, ae_override, av_override = dp_default.eval( + coord, + box, + atype, + fparam=fparam_override, + aparam=aparam_val, + atomic=True, + ) ref_path_default = os.path.join(base_dir, "fparam_aparam_default.expected") write_expected_ref( ref_path_default, @@ -186,6 +203,11 @@ def main(): "expected_f": f_d[0], "expected_v": av_d[0], }, + "override": { + "expected_e": ae_override[0, :, 0], + "expected_f": f_override[0], + "expected_v": av_override[0], + }, }, source_script="source/tests/infer/gen_fparam_aparam.py", ) From 43771fe024bb59b7f4dac43de78f1d8969646df2 Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Sat, 1 Aug 2026 22:18:10 +0800 Subject: [PATCH 2/2] fix(jax): complete fparam input coverage Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- source/api_cc/src/DeepPotJAX.cc | 5 +++-- .../api_cc/tests/deeppot_universal_test_common.h | 5 +++++ source/api_cc/tests/test_deeppot_universal.cc | 16 ++++++++++++++++ source/tests/infer/gen_fparam_aparam.py | 6 +++++- 4 files changed, 29 insertions(+), 3 deletions(-) diff --git a/source/api_cc/src/DeepPotJAX.cc b/source/api_cc/src/DeepPotJAX.cc index be9b3f21da..b203e41045 100644 --- a/source/api_cc/src/DeepPotJAX.cc +++ b/source/api_cc/src/DeepPotJAX.cc @@ -546,8 +546,9 @@ inline std::vector make_fparam_input( const std::vector& default_fparam) { if (dfparam == 0) { if (!fparam.empty()) { - throw deepmd::deepmd_exception( - "fparam was provided, but this model has dim_fparam=0."); + std::cerr << "WARNING: fparam was provided, but this model has " + "dim_fparam=0. The provided fparam will be ignored." + << std::endl; } return {}; } diff --git a/source/api_cc/tests/deeppot_universal_test_common.h b/source/api_cc/tests/deeppot_universal_test_common.h index 29be32acd1..ccf5906f73 100644 --- a/source/api_cc/tests/deeppot_universal_test_common.h +++ b/source/api_cc/tests/deeppot_universal_test_common.h @@ -148,6 +148,11 @@ inline std::vector fparam_aparam_cases() { "../../tests/infer/fparam_aparam.expected", 1e-7, 1e-4, /*supports_float=*/true}, {"pytorch_pt2", Backend::PTExpt, "../../tests/infer/fparam_aparam.pt2", + /*convert_pbtxt=*/false, nullptr, + "../../tests/infer/fparam_aparam.expected", 1e-7, 1e-4, + /*supports_float=*/true}, + {"jax_savedmodel", Backend::JAX, + "../../tests/infer/fparam_aparam.savedmodel", /*convert_pbtxt=*/false, nullptr, "../../tests/infer/fparam_aparam.expected", 1e-7, 1e-4, /*supports_float=*/true}}; diff --git a/source/api_cc/tests/test_deeppot_universal.cc b/source/api_cc/tests/test_deeppot_universal.cc index ca36b74c34..389e137c0a 100644 --- a/source/api_cc/tests/test_deeppot_universal.cc +++ b/source/api_cc/tests/test_deeppot_universal.cc @@ -1897,6 +1897,22 @@ TEST_P(FParamAParamDeepPotTest, Metadata) { EXPECT_FALSE(dp.has_default_fparam()); } +TEST_P(FParamAParamDeepPotTest, JAXRejectsMissingRequiredFParam) { + if (GetParam().backend != Backend::JAX) { + GTEST_SKIP() << "This regression targets JAX SavedModel tensor inputs."; + } + const std::vector coord = deepmd_test::deeppot_coord(); + const std::vector atype = deepmd_test::fparam_aparam_atype(); + const std::vector box = deepmd_test::deeppot_box(); + const std::vector aparam = deepmd_test::aparam_value(); + double energy = 0.0; + std::vector force, virial; + + EXPECT_THROW(dp.compute(energy, force, virial, coord, atype, box, + std::vector{}, aparam), + deepmd::deepmd_exception); +} + TEST_P(FParamAParamDeepPotTest, ComputeDouble) { check_fparam_compute_atomic(dp, *ref, GetParam().double_tol, deepmd_test::fparam_value()); diff --git a/source/tests/infer/gen_fparam_aparam.py b/source/tests/infer/gen_fparam_aparam.py index a530642de8..90594b29f5 100644 --- a/source/tests/infer/gen_fparam_aparam.py +++ b/source/tests/infer/gen_fparam_aparam.py @@ -80,7 +80,7 @@ def main(): print(f"Exporting to {savedmodel_default_path} ...") # noqa: T201 jax_deserialize_to_file(savedmodel_default_path, copy.deepcopy(data_default)) - # ---- 3. Export fparam_aparam.pt2 and .pth (without default_fparam) ---- + # ---- 3. Export fparam_aparam artifacts (without default_fparam) ---- config_no_default = copy.deepcopy(config) config_no_default["fitting_net"].pop("default_fparam", None) model_no_default = get_model(config_no_default) @@ -99,6 +99,10 @@ def main(): pt2_path, copy.deepcopy(data_no_default), do_atomic_virial=True ) + savedmodel_path = os.path.join(base_dir, "fparam_aparam.savedmodel") + print(f"Exporting to {savedmodel_path} ...") # noqa: T201 + jax_deserialize_to_file(savedmodel_path, copy.deepcopy(data_no_default)) + pth_path = os.path.join(base_dir, "fparam_aparam.pth") pth_exported = False print(f"Exporting to {pth_path} ...") # noqa: T201