Update type for kfac_jax Optimizer.State.

The forward declaration was recently removed; instead use the public API.

PiperOrigin-RevId: 624106473
Change-Id: Ia2043491c70b31c89d5025401b0f2d7ad36a972e

ferminet/train.py

@@ -127,7 +127,7 @@ def init_electrons(  # pylint: disable=dangerous-default-value
 
 
 # All optimizer states (KFAC and optax-based).
-OptimizerState = Union[optax.OptState, kfac_jax.optimizer.OptimizerState]
+OptimizerState = Union[optax.OptState, kfac_jax.Optimizer.State]
 OptUpdateResults = Tuple[networks.ParamTree, Optional[OptimizerState],
                          jnp.ndarray,
                          Optional[qmc_loss_functions.AuxiliaryLossData]]
@@ -331,7 +331,7 @@ def make_kfac_training_step(
   def step(
       data: networks.FermiNetData,
       params: networks.ParamTree,
-      state: kfac_jax.optimizer.OptimizerState,
+      state: kfac_jax.Optimizer.State,
       key: chex.PRNGKey,
       mcmc_width: jnp.ndarray,
   ) -> StepResults: