Fix https://github.com/google-research/t5x/issues/250
Fix https://github.com/google-research/t5x/issues/250

This change updates the `step` value of the `InferenceState` to always be `None` when returned from the `as_logical_axes` method. Originally, when using the `self.step` value the step leaf in the pytree would be a `jax.ShapeDtypeStruct` (when the partitioning first happens after the call to `jax.eval_shape` and cause an error because it was not iterable. If you get past that error (for example by checking for this object and replacing it with `None`, you would run into a similar error where the step is an interger loaded from the checkpoint, resulting in a similar not iterable error.

This fix sets this value to always be `None` so that later when the logical axis version of the `InferenceState` is passed to a `jax.tree_map` to be processed by the `flax.linen.partitioning.logical_to_mesh_axis` function, `jax.tree_map` short-circuiting on `None` valued leaves will skip processing the state. A possible alternative solution would be to remove the step value when setting up partitioning in-case there are any other situtations where we would want the step value in the logical axis tree, but I'm not sure if that would work and can't think of a reason that it would be needed.

The `InferenceState` still works as normal when not using the return value from the `.as_logical_axes` method, so when the actual computation is done, the true step value is present and the metrics are logged with the correct step.

PiperOrigin-RevId: 428582385

google-research/t5x