fix(jax): materialize default fparam in cxx api (deepmodeling#5849)

Closes deepmodeling#5658.

## Summary

- load and validate `get_default_fparam` when initializing a JAX
SavedModel in the C++ API
- materialize either caller-provided or stored frame parameters,
including one-frame-to-multiframe broadcasting, in both direct and
neighbor-list paths
- reject inconsistent tensor data/shape combinations, cast tensor rank
explicitly, and report TensorFlow C API allocation failure before
accessing tensor storage
- generate a JAX default-fparam fixture and extend the shared C++
regression matrix with default, distinct explicit override,
float/double, neighbor-list, multiframe, and invalid-size coverage

## Why existing tests missed this

The shared C++ default-fparam suite only included TorchScript and PT2
artifacts. Existing JAX C++ fixtures either had `dim_fparam == 0` or
supplied frame parameters explicitly, while Python JAX inference
materializes the saved default before invoking the model. Serialization
coverage therefore proved that the flag/getter were exported, but never
exercised C++ consumption of an omitted JAX `fparam`.

The explicit regression now uses `fparam=[0.5]`, distinct from the
stored default `[0.25852028]`, so it also detects an implementation that
ignores caller overrides.

## Validation

- `ruff format .`
- `ruff check .`
- `clang-format --dry-run --Werror` on all changed C++ files
- TensorFlow/JAX-only C++ build of `runUnitTests_cc` and
`deepmd_backend_jax`; both targets rebuilt successfully after the
allocation-safety follow-up
- 11 focused `DefaultFParamDeepPotTest` JAX SavedModel cases passed,
covering direct and neighbor-list inference, float/double, stored
defaults, distinct explicit overrides, two-frame broadcasting, metadata,
and invalid sizes

Coding agent: Codex
Codex version: codex-cli 0.144.4
Model: gpt-5.6-sol
Reasoning effort: xhigh


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **Bug Fixes**
* Hardened tensor creation with shape validation, overflow checks, and
safer zero-byte handling.
* Improved JAX frame-parameter behavior, including embedded model
defaults, stricter validation, and clear warnings when defaults are
unavailable.
* Added support for supplying per-frame or full-length frame-parameter
inputs, with broadcasting across frames.
* Ensures invalid frame-parameter sizes are rejected with better error
reporting.
* **Tests**
* Expanded JAX SavedModel test coverage for overrides, broadcasting, and
invalid input sizes.
* Updated/added generated JAX SavedModel artifacts and corresponding
expected reference outputs.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: njzjz-bot <njzjz.bot@gmail.com>
Co-authored-by: njzjz-bot <njzjz-bot@users.noreply.github.com>

Author: njzjz-bot

source/api_cc/include/DeepPotJAX.h

@@ -276,6 +276,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,60 @@ 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()) {
+      std::cerr << "WARNING: fparam was provided, but this model has "
+                   "dim_fparam=0. The provided fparam will be ignored."
+                << std::endl;
+    }
+    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,
@@ -675,6 +758,27 @@ 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
@@ -779,6 +883,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_);
@@ -942,6 +1048,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/deeppot_universal_test_common.h

@@ -148,6 +148,11 @@ inline std::vector<FParamAParamCase> 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}};

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"
@@ -359,6 +360,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}};
 }
 
@@ -532,6 +537,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();
@@ -543,6 +550,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);
   }
@@ -1918,6 +1926,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<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();
+  double energy = 0.0;
+  std::vector<double> force, virial;
+
+  EXPECT_THROW(dp.compute(energy, force, virial, coord, atype, box,
+                          std::vector<double>{}, aparam),
+               deepmd::deepmd_exception);
+}
+
 TEST_P(FParamAParamDeepPotTest, ComputeDouble) {
   check_fparam_compute_atomic<double>(dp, *ref, GetParam().double_tol,
                                       deepmd_test::fparam_value());
@@ -2007,17 +2031,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) {
@@ -2034,6 +2058,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,