fix(jax): materialize default fparam in cxx api

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<double> 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

source/api_cc/src/DeepPotJAX.cc

@@ -12,6 +12,7 @@
 #include <cstdio>
 #include <cstring>
 #include <iostream>
+#include <limits>
 #include <numeric>
 #include <ostream>
 #include <stdexcept>
@@ -439,10 +440,38 @@ inline std::vector<std::string> get_vector_string(
 template <typename T>
 inline TF_Tensor* create_tensor(const std::vector<T>& data,
                                 const std::vector<int64_t>& 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<size_t>(dim);
+    if (size_dim != 0 &&
+        element_count > std::numeric_limits<size_t>::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<size_t>::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<int>(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<double> make_charge_spin_input(
       std::to_string(nframes) + " frames).");
 }
 
+inline std::vector<double> make_fparam_input(
+    const std::vector<double>& fparam,
+    const int dfparam,
+    const int nframes,
+    const bool has_default_fparam,
+    const std::vector<double>& 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<double>* 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<size_t>(dfparam);
+  if (static_cast<size_t>(nframes) > std::numeric_limits<size_t>::max() / dim) {
+    throw deepmd::deepmd_exception("fparam element count overflows size_t.");
+  }
+  const size_t expected = static_cast<size_t>(nframes) * dim;
+  if (source->size() == expected) {
+    return *source;
+  }
+  if (source->size() == dim) {
+    std::vector<double> result(expected);
+    for (int ff = 0; ff < nframes; ++ff) {
+      std::copy(source->begin(), source->end(),
+                result.begin() + static_cast<size_t>(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 <typename T>
 inline void tensor_to_vector(std::vector<T>& 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<double>(ctx, "get_default_fparam",
+                                           func_vector, device, status);
+      if (static_cast<int>(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<ENERGYTYPE>& ener,
   std::vector<double> coord_double(coord.begin(), coord.end());
   std::vector<double> box_double(box.begin(), box.end());
   std::vector<double> fparam_double(fparam.begin(), fparam.end());
+  fparam_double = make_fparam_input(fparam_double, dfparam, nframes,
+                                    has_default_fparam_, default_fparam_);
   std::vector<double> aparam_double(aparam.begin(), aparam.end());
   std::vector<double> charge_spin_double =
       make_charge_spin_input(charge_spin, dchgspin, nframes, default_chg_spin_);
@@ -912,6 +1016,8 @@ void deepmd::DeepPotJAX::compute(std::vector<ENERGYTYPE>& ener,
   // float model interface
   std::vector<double> coord_double(coord.begin(), coord.end());
   std::vector<double> fparam_double(fparam.begin(), fparam.end());
+  fparam_double = make_fparam_input(fparam_double, dfparam, nframes,
+                                    has_default_fparam_, default_fparam_);
   std::vector<double> aparam_double(aparam.begin(), aparam.end());
   std::vector<double> charge_spin_double =
       make_charge_spin_input(charge_spin, dchgspin, nframes, default_chg_spin_);

source/api_cc/tests/test_deeppot_universal.cc

@@ -7,6 +7,7 @@
 #include <exception>
 #include <string>
 #include <type_traits>
+#include <utility>
 #include <vector>
 
 #include "DeepPot.h"
@@ -307,6 +308,10 @@ std::vector<DefaultFParamCase> 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<double> 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,17 +1950,17 @@ TEST_P(DefaultFParamDeepPotTest, ComputeWithEmptyFParamFloat) {
 }
 
 TEST_P(DefaultFParamDeepPotTest, ComputeWithExplicitFParamDouble) {
-  check_fparam_compute_simple<double>(dp, ref, GetParam().double_tol,
-                                      deepmd_test::fparam_value());
+  check_fparam_compute_simple<double>(dp, override_ref, GetParam().double_tol,
+                                      override_fparam);
 }
 
 TEST_P(DefaultFParamDeepPotTest, ComputeWithExplicitFParamFloat) {
   if (!GetParam().supports_float) {
     GTEST_SKIP() << backend_name(GetParam().backend)
                  << " does not provide float inference coverage.";
   }
-  check_fparam_compute_simple<float>(dp, ref, GetParam().float_tol,
-                                     deepmd_test::fparam_value());
+  check_fparam_compute_simple<float>(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<double>(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<float>(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<double> coord =
+      repeat_values(deepmd_test::deeppot_coord(), nframes);
+  const std::vector<int> atype = deepmd_test::fparam_aparam_atype();
+  const std::vector<double> box =
+      repeat_values(deepmd_test::deeppot_box(), nframes);
+  const std::vector<double> aparam =
+      repeat_values(deepmd_test::aparam_value(), nframes);
+  const std::vector<
+      std::pair<std::vector<double>, 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<double> expected_virial =
+        deepmd_test::total_virial(*expected_ref);
+    std::vector<double> energy, force, virial;
+    dp.compute(energy, force, virial, coord, atype, box, fparam, aparam);
+
+    ASSERT_EQ(energy.size(), static_cast<size_t>(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<size_t>(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<size_t>(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<double> coord = deepmd_test::deeppot_coord();
+  const std::vector<int> atype = deepmd_test::fparam_aparam_atype();
+  const std::vector<double> box = deepmd_test::deeppot_box();
+  const std::vector<double> aparam = deepmd_test::aparam_value();
+  const std::vector<double> invalid_fparam = {0.1, 0.2};
+  double energy = 0.0;
+  std::vector<double> force, virial;
+
+  EXPECT_THROW(dp.compute(energy, force, virial, coord, atype, box,
+                          invalid_fparam, aparam),
+               deepmd::deepmd_exception);
+}
+
 INSTANTIATE_TEST_SUITE_P(
     AvailableBackends,
     UniversalDeepPotTest,

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