Fix opt_state type mismatch when restoring checkpoints with Adam/LAMB optimizers

When using Adam or LAMB optimizers with reset_if_nan=True, the opt_state
loaded from checkpoint is a list but should be a tuple. This causes
jax.lax.cond at train.py:315 to fail during training. Convert opt_state_ckpt
to tuple when restoring from checkpoint to ensure type consistency.

Author: Qianruipku

ferminet/train.py

@@ -883,7 +883,8 @@ def learning_rate_schedule(t_: jnp.ndarray) -> jnp.ndarray:
   elif isinstance(optimizer, optax.GradientTransformation):
     # optax/optax-compatible optimizer (ADAM, LAMB, ...)
     opt_state = jax.pmap(optimizer.init)(params)
-    opt_state = opt_state_ckpt or opt_state  # avoid overwriting ckpted state
+    if opt_state_ckpt is not None:
+      opt_state = tuple(opt_state_ckpt)
     step = make_training_step(
         mcmc_step=mcmc_step,
         optimizer_step=make_opt_update_step(evaluate_loss, optimizer),