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", )