feat(jax): freeze models with hessian output

deepmd/jax/entrypoints/freeze.py

@@ -18,6 +18,7 @@ def freeze(
     *,
     checkpoint_folder: str,
     output: str,
+    hessian: bool = False,
     **kwargs: object,
 ) -> None:
     """Freeze a JAX checkpoint into a serialized model file.
@@ -30,6 +31,8 @@ def freeze(
     output : str
         Output model filename or prefix. The JAX model suffix is added when the
         filename has no supported backend suffix.
+    hessian : bool, default=False
+        Whether to include the Hessian in the frozen model outputs.
     **kwargs
         Other CLI arguments accepted for backend entry-point compatibility.
     """
@@ -46,4 +49,4 @@ def freeze(
         strict_prefer=True,
     )
     data = serialize_from_file(checkpoint_folder)
-    deserialize_to_file(output, data)
+    deserialize_to_file(output, data, hessian=hessian)

deepmd/jax/infer/deep_eval.py

@@ -318,6 +318,7 @@ def _get_request_defs(self, atomic: bool) -> list[OutputVariableDef]:
                     OutputVariableCategory.REDU,
                     OutputVariableCategory.DERV_R,
                     OutputVariableCategory.DERV_C_REDU,
+                    OutputVariableCategory.DERV_R_DERV_R,
                 )
             ]
 
@@ -470,6 +471,10 @@ def get_model_def_script(self) -> dict:
         """Get model definition script."""
         return json.loads(self.dp.get_model_def_script())
 
+    def get_has_hessian(self) -> bool:
+        """Check if the model has Hessian output."""
+        return self.get_model_def_script().get("hessian_mode", False)
+
     def get_model(self) -> Any:
         """Get the JAX model as BaseModel.
 

deepmd/jax/jax2tf/serialization.py

@@ -29,16 +29,22 @@
     BaseModel,
 )
 from deepmd.jax.utils.serialization import (
+    _prepare_hessian_model_def_script,
     _set_model_min_nbor_dist_from_data,
 )
 
 
-def deserialize_to_file(model_file: str, data: dict) -> None:
+def deserialize_to_file(model_file: str, data: dict, hessian: bool = False) -> None:
     """Deserialize the dictionary to a JAX/jax2tf SavedModel."""
     if model_file.endswith(".savedmodel"):
         model = BaseModel.deserialize(data["model"])
         _set_model_min_nbor_dist_from_data(model, data)
-        model_def_script = data["model_def_script"]
+        model_def_script, hessian = _prepare_hessian_model_def_script(
+            data["model_def_script"],
+            hessian,
+        )
+        if hessian:
+            model.enable_hessian()
         call_lower = model.call_common_lower
 
         tf_model = tf.Module()

deepmd/jax/model/hlo.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+import json
 from typing import (
     Any,
 )
@@ -30,6 +31,14 @@
         r_differentiable=True,
         c_differentiable=True,
     ),
+    "energy_hessian": OutputVariableDef(
+        "energy",
+        shape=[1],
+        reducible=True,
+        r_differentiable=True,
+        c_differentiable=True,
+        r_hessian=True,
+    ),
     "dos": OutputVariableDef(
         "dos",
         shape=[-1],
@@ -193,6 +202,10 @@ def model_output_def(self) -> ModelOutputDef:
         )
 
     def _output_var_def(self, name: str) -> OutputVariableDef:
+        if name == "energy" and json.loads(self.model_def_script).get(
+            "hessian_mode", False
+        ):
+            return OUTPUT_DEFS["energy_hessian"]
         if name in OUTPUT_DEFS:
             return OUTPUT_DEFS[name]
         # property models carry a user-defined output name (``var_name``) that

deepmd/jax/utils/serialization.py

@@ -185,7 +185,19 @@ def _check_compressed_hlo_exportable(data: dict) -> None:
         )
 
 
-def deserialize_to_file(model_file: str, data: dict) -> None:
+def _prepare_hessian_model_def_script(
+    model_def_script: dict,
+    hessian: bool,
+) -> tuple[dict, bool]:
+    """Return a copied model definition and whether Hessian should be enabled."""
+    model_def_script = model_def_script.copy()
+    hessian = hessian or model_def_script.get("hessian_mode", False)
+    if hessian:
+        model_def_script["hessian_mode"] = True
+    return model_def_script, hessian
+
+
+def deserialize_to_file(model_file: str, data: dict, hessian: bool = False) -> None:
     """Deserialize the dictionary to a model file.
 
     Parameters
@@ -194,9 +206,14 @@ def deserialize_to_file(model_file: str, data: dict) -> None:
         The model file to be saved.
     data : dict
         The dictionary to be deserialized.
+    hessian : bool, default=False
+        Whether to include the Hessian in the model outputs.
     """
     if model_file.endswith(".jax"):
-        model_def_script = data["model_def_script"].copy()
+        model_def_script, hessian = _prepare_hessian_model_def_script(
+            data["model_def_script"],
+            hessian,
+        )
         min_nbor_dist = _to_optional_float(data.get("min_nbor_dist"))
         if min_nbor_dist is None:
             min_nbor_dist = _to_optional_float(
@@ -209,6 +226,9 @@ def deserialize_to_file(model_file: str, data: dict) -> None:
                 model_key: BaseModel.deserialize(data["model"]["model_dict"][model_key])
                 for model_key in model_def_script["model_dict"]
             }
+            if hessian:
+                for model in models.values():
+                    model.enable_hessian()
             state = {
                 "models": {
                     model_key: nnx.split(model)[1].to_pure_dict()
@@ -217,6 +237,8 @@ def deserialize_to_file(model_file: str, data: dict) -> None:
             }
         else:
             model = BaseModel.deserialize(data["model"])
+            if hessian:
+                model.enable_hessian()
             _, state = nnx.split(model)
             state = state.to_pure_dict()
         with ocp.Checkpointer(
@@ -233,7 +255,12 @@ def deserialize_to_file(model_file: str, data: dict) -> None:
         _check_compressed_hlo_exportable(data)
         model = BaseModel.deserialize(data["model"])
         _set_model_min_nbor_dist_from_data(model, data)
-        model_def_script = data["model_def_script"]
+        model_def_script, hessian = _prepare_hessian_model_def_script(
+            data["model_def_script"],
+            hessian,
+        )
+        if hessian:
+            model.enable_hessian()
         call_lower = model.call_common_lower
 
         nf, nloc, nghost = jax_export.symbolic_shape("nf, nloc, nghost")
@@ -298,6 +325,7 @@ def call_lower_with_fixed_do_atomic_virial(
         serialized_atomic_virial_no_ghost = exported_atomic_virial_no_ghost.serialize()
 
         data = data.copy()
+        data["model_def_script"] = model_def_script
         data.setdefault("@variables", {})
         data["@variables"]["stablehlo"] = np.void(serialized)
         data["@variables"]["stablehlo_atomic_virial"] = np.void(
@@ -344,7 +372,7 @@ def call_lower_with_fixed_do_atomic_virial(
             deserialize_to_file as deserialize_to_savedmodel,
         )
 
-        return deserialize_to_savedmodel(model_file, data)
+        deserialize_to_savedmodel(model_file, data, hessian=hessian)
     else:
         raise ValueError("Unsupported file extension")
 

deepmd/main.py

@@ -350,6 +350,12 @@ def main_parser() -> argparse.ArgumentParser:
         type=str,
         help="(Supported backend: PyTorch) Task head (alias: model branch) to freeze if in multi-task mode.",
     )
+    parser_frz.add_argument(
+        "--hessian",
+        action="store_true",
+        default=False,
+        help="(Supported backend: JAX) Add the Hessian to the frozen model output.",
+    )
     parser_frz.add_argument(
         "--lower-kind",
         default="nlist",

source/tests/jax/test_training.py

@@ -24,6 +24,9 @@
 import numpy as np
 import optax
 
+from deepmd.dpmodel.output_def import (
+    OutputVariableCategory,
+)
 from deepmd.dpmodel.train import (
     TrainEntrypointOptions,
 )
@@ -40,6 +43,12 @@
 from deepmd.jax.env import (
     jnp,
 )
+from deepmd.jax.infer.deep_eval import (
+    DeepEval,
+)
+from deepmd.jax.model.hlo import (
+    HLO,
+)
 from deepmd.jax.train.trainer import (
     DPTrainer,
     _copy_matching_state_tree,
@@ -527,17 +536,19 @@ def test_update_sel_supports_multitask(self, get_nbor_stat, get_data) -> None:
     def test_freeze_entrypoint_uses_checkpoint_pointer(
         self, serialize_from_file, deserialize_to_file
     ) -> None:
-        """Freeze resolves the stable checkpoint pointer without Hessian options."""
+        """Freeze resolves the stable checkpoint pointer and forwards Hessian."""
         checkpoint_dir = self.work_dir / "ckpt"
         checkpoint_dir.mkdir()
         (checkpoint_dir / "checkpoint").write_text("model-1.jax")
         serialize_from_file.return_value = {"model": {}, "model_def_script": {}}
 
-        freeze(checkpoint_folder=str(checkpoint_dir), output="frozen_model")
+        freeze(
+            checkpoint_folder=str(checkpoint_dir), output="frozen_model", hessian=True
+        )
 
         serialize_from_file.assert_called_once_with(str(checkpoint_dir / "model-1.jax"))
         deserialize_to_file.assert_called_once_with(
-            "frozen_model.hlo", serialize_from_file.return_value
+            "frozen_model.hlo", serialize_from_file.return_value, hessian=True
         )
 
     @patch("deepmd.jax.entrypoints.main.freeze")
@@ -549,11 +560,44 @@ def test_main_dispatches_freeze(self, freeze_entrypoint) -> None:
             log_path=None,
             checkpoint_folder=".",
             output="frozen_model",
+            hessian=False,
         )
 
         main(args)
 
         freeze_entrypoint.assert_called_once()
+        self.assertIn("hessian", freeze_entrypoint.call_args.kwargs)
+        self.assertFalse(freeze_entrypoint.call_args.kwargs["hessian"])
+
+    def test_hlo_hessian_mode_updates_output_def(self) -> None:
+        """HLO output definition should expose Hessian when requested."""
+        hlo = object.__new__(HLO)
+        hlo._model_output_type = ["energy"]
+        hlo.model_def_script = json.dumps({"hessian_mode": True})
+
+        output_def = hlo.model_output_def()
+
+        self.assertTrue(output_def["energy"].r_hessian)
+        self.assertIn("energy_derv_r_derv_r", output_def.keys())
+
+    def test_deep_eval_requests_hessian_for_hessian_model(self) -> None:
+        """Non-atomic JAX evaluation should request Hessian outputs."""
+        hlo = object.__new__(HLO)
+        hlo._model_output_type = ["energy"]
+        hlo.model_def_script = json.dumps({"hessian_mode": True})
+        deep_eval = object.__new__(DeepEval)
+        deep_eval.output_def = hlo.model_output_def()
+        deep_eval.dp = SimpleNamespace(
+            get_model_def_script=lambda: json.dumps({"hessian_mode": True})
+        )
+
+        request_defs = deep_eval._get_request_defs(atomic=False)
+
+        self.assertTrue(deep_eval.get_has_hessian())
+        self.assertIn(
+            OutputVariableCategory.DERV_R_DERV_R,
+            {odef.category for odef in request_defs},
+        )
 
 
 def test_jax_finetune_state_copy_preserves_random_fitting_target_leaves() -> None: