jax: align training/model support with argcheck schema

## Summary

While comparing the unified input schema in `deepmd/utils/argcheck.py` with the JAX backend implementation, I found several schema-exposed options that are not implemented or not wired up in `deepmd/jax` yet.

This issue tracks the gaps so we can either implement them for JAX or make the schema/docs explicitly say that these options are unavailable in the JAX backend.

## Gaps

### 1. Non-energy fittings/models are registered, but JAX training only accepts `loss.type: ener`

Schema/backend registrations:

- `deepmd/utils/argcheck.py` exposes fitting types including `ener`, `dos`, `property`, `polar`, and `dipole`.
- The JAX backend registers matching fitting/model classes for `dos`, `property`, `dipole`, and `polar` under `deepmd/jax/fitting/fitting.py` and `deepmd/jax/model/*.py`.

Implementation:

- `deepmd/jax/train/trainer.py` `_build_losses()` only accepts `loss_type == "ener"`; any other loss raises `RuntimeError("unknown loss type ...")`.

Impact:

- JAX can construct non-energy model/fitting objects, but JAX training cannot train the corresponding DOS/property/tensor tasks because the loss dispatch rejects non-energy losses.

### 2. `model.spin` is accepted by the common schema, but JAX explicitly rejects spin models

Schema/docs:

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

Implementation:

- `deepmd/jax/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 JAX model construction time.

### 3. Learning-rate schema exposes `exp`, `cosine`, and `wsd`, but JAX only implements `exp`

Schema/docs:

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

Implementation:

- `deepmd/jax/train/trainer.py` `_get_lr_and_coef()` dispatches only `type == "exp"` to `LearningRateExp`; all other types raise `RuntimeError("unknown learning_rate type ...")`.

Impact:

- `learning_rate.type=cosine` or `learning_rate.type=wsd` is schema-exposed but unavailable in JAX training.

### 4. The optimizer schema is not wired into JAX training

Schema/docs:

- `deepmd/utils/argcheck.py` exposes an `optimizer` block and optimizer variants.

Implementation:

- `deepmd/jax/train/trainer.py` builds an Optax chain directly with `optax.scale_by_adam()` and `_scale_by_global_learning_rate()`; it does not read `config["optimizer"]` or dispatch optimizer variants/parameters.

Impact:

- Optimizer options accepted by the common input schema have no effect in JAX training. At minimum this should be documented as unsupported; ideally JAX should consume the optimizer config it can support.

### 5. Several model-level schema variants are not present in the JAX backend registry

Schema/docs:

- `deepmd/utils/argcheck.py` exposes model variants such as `frozen`, `pairtab`, and `linear_ener` in the common model schema.

Implementation observation:

- JAX model registrations are limited to the task models (`ener`, `dos`, `dipole`, `polar`, `property`) plus `zbl`; I did not find JAX `BaseModel.register(...)` implementations for `frozen`, top-level `pairtab`, or `linear_ener`.
- `model.type: dpa4` / `dpa4_ener` are marked PyTorch-only in the schema, so those are not counted here as JAX implementation gaps.

Impact:

- Users can discover these model variants from the unified schema, but JAX cannot construct them unless there is a separate compatibility path.

## Expected resolution

Either:

1. implement/wire these options in the JAX backend, or
2. update `argcheck.py`/docs to mark them as unsupported for JAX so schema-valid configurations do not fail later in backend-specific code.

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

## Implementation checklist

### Model, loss, and optimizer support

- [ ] Train non-energy DOS, property, dipole, and polar models with matching losses (#5865)
- [ ] Implement JAX spin models or reject them during backend-aware validation (#5866)
- [x] Support `learning_rate.type: cosine` and `learning_rate.type: wsd` (#5776)
- [ ] Consume the optimizer schema and support the applicable optimizer options (#5867; #5771 tracks DPA4 `HybridMuon`)
- [ ] Implement or explicitly reject/document `frozen`, `pairtab`, and `linear_ener` model variants (#5868)

### Training support

- [ ] Derive training steps from `training.numb_epoch` and its aliases (#5869)
- [ ] Support multi-task `training.num_epoch_dict` (#5869)
- [ ] Honor `training.stat_file` when computing or loading statistics (#5870)
- [ ] Implement `training.profiling` and `training.enable_profiler` (#5871)
- [ ] Implement `training.tensorboard*` output (#5872)
- [ ] Apply `training.mixed_precision` to training (#5873)
- [x] Apply `training.change_bias_after_training` after the final step (#5745)