test: align test doubles with the promoted capability contracts

DummyModel (jax2tf serialization test) gains get_var_name -> None; the
_DescriptorWithStats stub gains the stat-flag class defaults that
merge_env_stat now reads directly.

Author: Han Wang

source/jax2tf_tests/test_serialization.py

@@ -122,6 +122,10 @@ def has_default_chg_spin(self) -> bool:
         def get_default_chg_spin(self) -> None:
             return None
 
+        def get_var_name(self) -> None:
+            # non-property model, matching the make_base_model default
+            return None
+
     class DummyChargeSpinModel(DummyModel):
         dim_chg_spin = 2
 

source/tests/jax/test_training.py

@@ -496,6 +496,11 @@ def __call__(self) -> tuple[np.ndarray, np.ndarray]:
 
 
 class _DescriptorWithStats:
+    # stat-behavior flags merge_env_stat reads directly on any
+    # Descriptor/DescriptorBlock (class defaults on the real bases)
+    set_davg_zero = False
+    set_stddev_constant = False
+
     def __init__(self, stats: dict[str, StatItem]) -> None:
         self.stats = stats
         self.davg = np.asarray([0.0], dtype=np.float64)