fix(jax,tf2): deep-copy and propagate type_map in ZBL model factories

[wanghan-iapcm]

Problem

Fixes #5676. The JAX and TF2 ZBL model factories (get_zbl_model), unlike the standard-model factories and the dpmodel ZBL path, mutated the caller's config dict in place (popping the descriptor and fitting_net type) and did not inject type_map into the descriptor and fitting sub-configs. The in-place mutation is inconsistent with the neighboring factories, and the missing type_map leaves the descriptor and fitting without the model type map that standard/dpmodel construction provides.

Fix

Deep-copy the input and inject type_map into the descriptor and fitting configs in both get_zbl_model factories, mirroring the dpmodel ZBL path.

The issue also suspected the retained JAX checkpoint metadata could be corrupted by the in-place mutation. That does not actually happen: the JAX trainer passes a deepcopy of the model params into the factory, so the stored model_def_script is never mutated. No change is made for that sub-claim.

Test

Adds source/tests/jax/test_zbl_model.py and a TF2-gated source/tests/consistent/test_tf2_zbl_model.py, each asserting the factory leaves the input dict unchanged and that the constructed descriptor and fitting carry the model type_map. Both assertions fail on master (the input is mutated and the sub-config type_map is None) and pass with the fix. The TF2 test is gated on INSTALLED_TF2 (DEEPMD_TEST_TF2=1).

Summary by CodeRabbit

• Bug Fixes

  • Improved ZBL model setup so input configuration data is no longer modified during model creation.
  • Ensured type mappings are consistently applied to both descriptor and fitting settings when building ZBL models.

• Tests

  • Added coverage for JAX and TF2 ZBL model factories to verify configuration immutability and correct type mapping behavior.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Pro

Run ID: ca4a5679-ad9c-4d41-b616-b63f0b08568c

📥 Commits

Reviewing files that changed from the base of the PR and between bbc3908 and 4987e73.

📒 Files selected for processing (4)

• deepmd/jax/model/model.py
• deepmd/tf2/model/model.py
• source/tests/consistent/test_tf2_zbl_model.py
• source/tests/jax/test_zbl_model.py

---

📝 Walkthrough

Walkthrough

The JAX and TF2 get_zbl_model factory functions were modified to deep-copy input data before mutation and to reorder setup so type_map is injected into both descriptor and fitting_net configs prior to component construction. New unit tests were added for both backends verifying non-mutation and type_map propagation.

Changes

ZBL Model Factory Fixes

Layer / File(s)	Summary
JAX get_zbl_model non-mutation and type_map injection deepmd/jax/model/model.py, source/tests/jax/test_zbl_model.py	get_zbl_model deep-copies input data, derives descriptor["ntypes"] from type_map, assigns descriptor["type_map"] and fitting_net["type_map"] before building components; new tests verify no input mutation and correct type_map injection.
TF2 get_zbl_model non-mutation and type_map injection deepmd/tf2/model/model.py, source/tests/consistent/test_tf2_zbl_model.py	get_zbl_model deep-copies input data, updates descriptor ntypes/type_map, pops type fields, and sets fitting_net["type_map"]; new tests verify no input mutation and correct type_map propagation.

Estimated code review effort: 2 (Simple) | ~12 minutes

Sequence Diagram(s)

sequenceDiagram
  participant Caller
  participant get_zbl_model
  participant Descriptor
  participant FittingNet
  participant DPZBLModel

  Caller->>get_zbl_model: data (config dict)
  get_zbl_model->>get_zbl_model: deepcopy(data)
  get_zbl_model->>Descriptor: ntypes, type_map
  get_zbl_model->>Descriptor: construct(descriptor_kwargs)
  get_zbl_model->>FittingNet: type_map
  get_zbl_model->>FittingNet: construct(fitting_kwargs)
  get_zbl_model->>DPZBLModel: assemble(Descriptor, FittingNet)
  DPZBLModel-->>Caller: model instance

Loading

Related Issues: #5676

Suggested labels: bug, JAX, TF2, tests

Suggested reviewers: njzjz, iProzd

_{🎉 A tiny fix, deep and wise,
No more mutation in disguise.
Type maps flow where they belong,
JAX and TF2, both made strong.
Tests stand guard, the configs stay true. 🐇}

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly summarizes the main change: deep-copying ZBL model configs and propagating type_map for JAX and TF2.
Linked Issues check	✅ Passed	The PR matches the issue by deep-copying JAX/TF2 ZBL configs, avoiding caller mutation, and injecting type_map into descriptor and fitting configs.
Out of Scope Changes check	✅ Passed	The added tests and model edits are directly tied to the stated ZBL config and type_map fix.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 81.15%. Comparing base (dd38b35) to head (4987e73).
⚠️ Report is 1 commits behind head on master.

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #5725      +/-   ##
==========================================
- Coverage   81.26%   81.15%   -0.11%     
==========================================
  Files         988      988              
  Lines      110877   110889      +12     
  Branches     4234     4233       -1     
==========================================
- Hits        90103    89996     -107     
- Misses      19249    19369     +120     
+ Partials     1525     1524       -1     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.