Fix pytype issue when setting step=None in InferenceState. Also r…

…eplace `Optional` with `| None`.

PiperOrigin-RevId: 557479578

t5x/train_state.py

@@ -14,7 +14,7 @@
 
 """Train state for passing around objects during training."""
 
-from typing import Any, Mapping, MutableMapping, Optional, Tuple
+from typing import Any, Mapping, MutableMapping, Tuple
 
 from flax import traverse_util
 import flax.core
@@ -24,7 +24,6 @@
 import flax.struct
 import jax.numpy as jnp
 from t5x import optimizers
-
 import typing_extensions
 
 EMPTY_DICT = flax.core.freeze({})
@@ -117,11 +116,11 @@ class FlaxOptimTrainState(flax.struct.PyTreeNode):
   """Simple train state for holding parameters, step, optimizer state."""
   _optimizer: optimizers.OptimizerType
   # Contains axis metadata (e.g., names) matching parameter tree.
-  params_axes: Optional[FrozenVariableDict] = None
+  params_axes: FrozenVariableDict | None = None
   # Flax mutable fields.
   flax_mutables: FrozenDict = EMPTY_DICT
   # Contains axis metadata (e.g., names) matching flax_mutables tree.
-  flax_mutables_axes: Optional[FrozenVariableDict] = None
+  flax_mutables_axes: FrozenVariableDict | None = None
 
   @classmethod
   def create(
@@ -228,11 +227,11 @@ def as_logical_axes(self) -> 'FlaxOptimTrainState':
 class InferenceState(flax.struct.PyTreeNode):
   """State compatible with FlaxOptimTrainState without optimizer state."""
 
-  step: jnp.ndarray
+  step: jnp.ndarray | None
   params: flax_scope.FrozenVariableDict
-  params_axes: Optional[flax_scope.FrozenVariableDict] = None
+  params_axes: flax_scope.FrozenVariableDict | None = None
   flax_mutables: flax_scope.FrozenDict = EMPTY_DICT
-  flax_mutables_axes: Optional[flax_scope.FrozenVariableDict] = None
+  flax_mutables_axes: flax_scope.FrozenVariableDict | None = None
 
   @classmethod
   def create(cls, model_variables: FrozenVariableDict) -> 'InferenceState':