tf2: align backend support with argcheck schema

[njzjz-bot]

Summary

While comparing deepmd/utils/argcheck.py with the deepmd/tf2 backend, I found several common-schema options that are not implemented or not fully wired up in TF2.

Some options are documented as Supported Backend: TensorFlow; if those are intended to mean only the legacy deepmd/tf backend and not deepmd/tf2, the docs/schema should make that distinction explicit.

Gaps

1. Legacy TensorFlow-only descriptors are schema-exposed, but TF2 does not register them

Schema/docs:

• deepmd/utils/argcheck.py exposes TensorFlow-supported descriptor variants including loc_frame, se_a_tpe/se_a_ebd, se_a_ebd_v2/se_a_tpe_v2, and se_a_mask.

Implementation observation:

• deepmd/tf2/descriptor/ registers the common dpmodel-style descriptors (se_e2_a, se_e2_r, se_e3, se_e3_tebd, se_atten, se_atten_v2, dpa2, dpa3, hybrid), but not loc_frame, se_a_tpe/se_a_ebd, se_a_ebd_v2/se_a_tpe_v2, or se_a_mask.

Impact:

• A user reading Supported Backend: TensorFlow may reasonably expect these to work with the TF2 backend, but TF2 cannot construct them.

2. TF-only model features are schema-exposed, but TF2 does not implement them

Schema/docs:

• deepmd/utils/argcheck.py exposes TF-only model-level features such as type_embedding, modifier, compress, and hybrid model pairwise_dprc.
• It also exposes model variants such as frozen, pairtab, and linear_ener in the common model schema.

Implementation:

• deepmd/tf2/model/model.py explicitly rejects model-level type_embedding with a ValueError and does not implement the legacy TF model-level path.
• TF2 model registrations are limited to task models (ener, dos, dipole, polar, property) plus zbl; I did not find TF2 BaseModel.register(...) implementations for frozen, top-level pairtab, linear_ener, or pairwise_dprc.

Impact:

• Schema-visible TensorFlow/common model options are unavailable in TF2 or fail during model construction.

3. model.spin is accepted by the common schema, but TF2 explicitly rejects spin models

Schema/docs:

• deepmd/utils/argcheck.py exposes the top-level model.spin block.

Implementation:

• deepmd/tf2/model/model.py raises NotImplementedError("Spin model is not implemented yet.") when model.type == "standard" and spin is present.

Impact:

• Schema-valid spin model configs fail at TF2 model construction time.

4. Learning-rate schema exposes exp, cosine, and wsd, but TF2 hard-codes LearningRateExp

Schema/docs:

• deepmd/utils/argcheck.py registers learning_rate.type variants exp, cosine, and wsd.

Implementation:

• deepmd/tf2/train/trainer.py directly constructs self.lr_schedule = LearningRateExp(**lr_params) without dispatching on learning_rate.type.

Impact:

• learning_rate.type=cosine or learning_rate.type=wsd is schema-exposed but unavailable in TF2 training.

5. Some schema loss types are not implemented in TF2

Schema/docs:

• deepmd/utils/argcheck.py exposes loss variants including ener, dens, ener_spin, dos, population, property, and tensor.

Implementation:

• deepmd/tf2/train/trainer.py get_loss() only supports ener, dos, tensor, and property.
• There is no TF2 dispatch for dens, ener_spin, or population.

Impact:

• These loss types are visible in the common schema but fail in TF2 training. If they are intentionally limited to other backends, the schema/docs should spell out the backend restriction.

6. Optimizer support is narrower than the common schema suggests

Schema/docs:

• deepmd/utils/argcheck.py exposes optimizer variants Adam, AdamW, LKF, AdaMuon, and HybridMuon.

Implementation:

• deepmd/tf2/train/trainer.py _build_optimizer() supports Adam and AdamW only; other optimizer names raise ValueError("Unsupported optimizer type for tf2: ...").

Impact:

• The common optimizer schema is broader than TF2 support. If only Adam/AdamW are intended for TF2, the docs/schema should make that explicit.

Expected resolution

Either:

1. implement/wire these options in the TF2 backend, or
2. update argcheck.py/docs to clearly distinguish legacy tf vs tf2 support and mark unsupported TF2 options before users hit backend runtime errors.

Filed by OpenClaw 2026.6.11 (e085fa1), model: custom-chat-jinzhezeng-group/gpt-5.5.

Implementation checklist

Model, loss, and optimizer support

• Implement or explicitly reject/document legacy TensorFlow-only descriptors in TF2
• Implement or explicitly reject/document TF-only model features and model variants in TF2
• Implement TF2 spin models or reject them during backend-aware validation
• Support learning_rate.type: cosine and learning_rate.type: wsd (fix(train): dispatch registered learning rate schedules #5776)
• Implement or explicitly reject/document dens, ener_spin, and population losses
• Support or explicitly restrict schema-exposed optimizers beyond Adam and AdamW (Add HybridMuon optimizer support to JAX and TF2 for DPA4 training #5771 tracks DPA4 HybridMuon)

Training support

• Derive training steps from training.numb_epoch and its aliases
• Support multi-task training.num_epoch_dict
• Implement training.profiling and training.enable_profiler
• Implement training.mixed_precision
• Apply training.change_bias_after_training after the final step (fix(train): share post-training bias adjustment #5745)

[njzjz-bot]

Addendum: training option coverage

I re-checked the training_data_args / validation_data_args / mixed_precision_args / training_args path against TF2 on current master (f73de32). Additional training-option gaps / clarifications:

• training.numb_epoch / aliases and multi-task training.num_epoch_dict: the schema allows epoch-based training, and legacy deepmd/tf/entrypoints/train.py has logic to compute numb_steps from numb_epoch and total_numb_batch. In contrast, deepmd/tf2/train/trainer.py still initializes self.num_steps = int(training_params["numb_steps"]); I did not find equivalent epoch -> step conversion or num_epoch_dict handling in the TF2 path.
• training.profiling and training.enable_profiler: these are schema-visible, but TF2 currently raises NotImplementedError for both in _validate_unsupported_config().
• training.mixed_precision: this is schema-visible, but TF2 currently raises NotImplementedError("TF2 training does not support mixed_precision yet.") when it is set.

Clarification: I did not count training.stat_file, training.tensorboard*, or training.change_bias_after_training as TF2 gaps in this addendum, because current TF2 code appears to wire those through (stat_file_path construction/passing, TensorBoard writer hooks, and a post-training _change_bias_after_training() branch).

Addendum by OpenClaw 2026.6.11 (e085fa1), model: custom-chat-jinzhezeng-group/gpt-5.5.