Repository: google/flax

Issue:
Inconsistent type annotation of `step` in `training.checkpoints`
### System information
- OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Any
- Flax, jax, jaxlib versions (obtain with `pip show flax jax jaxlib`: 
```
Name: flax
Version: 0.5.2
Summary: Flax: A neural network library for JAX designed for flexibility
Home-page: https://github.com/google/flax
Author: Flax team
Author-email: flax-dev@google.com
License: UNKNOWN
Location: /home/qys/Research/generative-distribution-shift/.venv/lib/python3.8/site-packages
Requires: jax, matplotlib, msgpack, numpy, optax, PyYAML, rich, typing-extensions
Required-by:
---
Name: jax
Version: 0.3.14
Summary: Differentiate, compile, and transform Numpy code.
Home-page: https://github.com/google/jax
Author: JAX team
Author-email: jax-dev@google.com
License: Apache-2.0
Location: /home/qys/Research/generative-distribution-shift/.venv/lib/python3.8/site-packages
Requires: absl-py, etils, numpy, opt-einsum, scipy, typing-extensions
Required-by: chex, flax, optax
---
Name: jaxlib
Version: 0.3.14
Summary: XLA library for JAX
Home-page: https://github.com/google/jax
Author: JAX team
Author-email: jax-dev@google.com
License: Apache-2.0
Location: /home/qys/Research/generative-distribution-shift/.venv/lib/python3.8/site-packages
Requires: absl-py, flatbuffers, numpy, scipy
Required-by: chex, optax
```
- Python version: Python 3.8.10
- GPU/TPU model and memory: Any
- CUDA version (if applicable): N/A


### Problem you have encountered:

The argument `step` has inconsistent type annotation. For example, the docstring of `save_checkpoint` says it can be either an integer or a float

https://github.com/google/flax/blob/0740ef63c4eae05de58d80f85a05fc23bb8b3261/flax/training/checkpoints.py#L241

but the docstring of `restore_checkpoint` says it must be an integer

https://github.com/google/flax/blob/0740ef63c4eae05de58d80f85a05fc23bb8b3261/flax/training/checkpoints.py#L403-L404

However, the example given by `restore_checkpoint` hints that a float is fine

https://github.com/google/flax/blob/0740ef63c4eae05de58d80f85a05fc23bb8b3261/flax/training/checkpoints.py#L390-L397

The documentation is also inconsistent with the actual type annotation. This makes linters like mypy and pyright unhappy.

https://github.com/google/flax/blob/0740ef63c4eae05de58d80f85a05fc23bb8b3261/flax/training/checkpoints.py#L225

https://github.com/google/flax/blob/0740ef63c4eae05de58d80f85a05fc23bb8b3261/flax/training/checkpoints.py#L384

### What you expected to happen:

The correct type annotation should be `step: Union[int, float]`.

### Logs, error messages, etc:

N/A

### Steps to reproduce:

N/A
Inconsistent type annotation of `step` in `training.checkpoints`
### System information
- OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Any
- Flax, jax, jaxlib versions (obtain with `pip show flax jax jaxlib`: 
```
Name: flax
Version: 0.5.2
Summary: Flax: A neural network library for JAX designed for flexibility
Home-page: https://github.com/google/flax
Author: Flax team
Author-email: flax-dev@google.com
License: UNKNOWN
Location: /home/qys/Research/generative-distribution-shift/.venv/lib/python3.8/site-packages
Requires: jax, matplotlib, msgpack, numpy, optax, PyYAML, rich, typing-extensions
Required-by:
---
Name: jax
Version: 0.3.14
Summary: Differentiate, compile, and transform Numpy code.
Home-page: https://github.com/google/jax
Author: JAX team
Author-email: jax-dev@google.com
License: Apache-2.0
Location: /home/qys/Research/generative-distribution-shift/.venv/lib/python3.8/site-packages
Requires: absl-py, etils, numpy, opt-einsum, scipy, typing-extensions
Required-by: chex, flax, optax
---
Name: jaxlib
Version: 0.3.14
Summary: XLA library for JAX
Home-page: https://github.com/google/jax
Author: JAX team
Author-email: jax-dev@google.com
License: Apache-2.0
Location: /home/qys/Research/generative-distribution-shift/.venv/lib/python3.8/site-packages
Requires: absl-py, flatbuffers, numpy, scipy
Required-by: chex, optax
```
- Python version: Python 3.8.10
- GPU/TPU model and memory: Any
- CUDA version (if applicable): N/A


### Problem you have encountered:

The argument `step` has inconsistent type annotation. For example, the docstring of `save_checkpoint` says it can be either an integer or a float

https://github.com/google/flax/blob/0740ef63c4eae05de58d80f85a05fc23bb8b3261/flax/training/checkpoints.py#L241

but the docstring of `restore_checkpoint` says it must be an integer

https://github.com/google/flax/blob/0740ef63c4eae05de58d80f85a05fc23bb8b3261/flax/training/checkpoints.py#L403-L404

However, the example given by `restore_checkpoint` hints that a float is fine

https://github.com/google/flax/blob/0740ef63c4eae05de58d80f85a05fc23bb8b3261/flax/training/checkpoints.py#L390-L397

The documentation is also inconsistent with the actual type annotation. This makes linters like mypy and pyright unhappy.

https://github.com/google/flax/blob/0740ef63c4eae05de58d80f85a05fc23bb8b3261/flax/training/checkpoints.py#L225

https://github.com/google/flax/blob/0740ef63c4eae05de58d80f85a05fc23bb8b3261/flax/training/checkpoints.py#L384

### What you expected to happen:

The correct type annotation should be `step: Union[int, float]`.

### Logs, error messages, etc:

N/A

### Steps to reproduce:

N/A


Discussion:

Resolving patch:
diff --git a/flax/training/checkpoints.py b/flax/training/checkpoints.py
--- a/flax/training/checkpoints.py
+++ b/flax/training/checkpoints.py
@@ -58,7 +58,7 @@
 
 
 def _checkpoint_path(ckpt_dir: str,
-                     step: Union[int, str],
+                     step: Union[int, float, str],
                      prefix: str = 'checkpoint_') -> str:
   return os.path.join(ckpt_dir, f'{prefix}{step}')
 
@@ -113,7 +113,7 @@ def _save_gdas(gda_manager: GlobalAsyncCheckpointManager,
 def _restore_gdas(state_dict,
                   target: Optional[Any],
                   ckpt_path: str,
-                  step: Optional[int] = None,
+                  step: Optional[Union[int, float]] = None,
                   gda_manager: Optional[GlobalAsyncCheckpointManager] = None):
 
   # When target is a single leaf instead of a pytree dict.
@@ -222,7 +222,7 @@ def save_async(self, task: Callable[[], Any]):
 
 def save_checkpoint(ckpt_dir: Union[str, os.PathLike],
                     target: PyTree,
-                    step: int,
+                    step: Union[int, float],
                     prefix: str = 'checkpoint_',
                     keep: int = 1,
                     overwrite: bool = False,
@@ -381,7 +381,7 @@ def latest_checkpoint(ckpt_dir: Union[str, os.PathLike],
 def restore_checkpoint(
     ckpt_dir: Union[str, os.PathLike],
     target: Optional[Any],
-    step: Optional[int] = None,
+    step: Optional[Union[int, float]] = None,
     prefix: str = 'checkpoint_',
     parallel: bool = True,
     gda_manager: Optional[GlobalAsyncCheckpointManager] = None) -> PyTree:
@@ -400,7 +400,7 @@ def restore_checkpoint(
     ckpt_dir: str: checkpoint file or directory of checkpoints to restore from.
     target: matching object to rebuild via deserialized state-dict. If None, the
       deserialized state-dict is returned as-is.
-    step: int: step number to load or None to load latest. If specified,
+    step: int or float: step number to load or None to load latest. If specified,
       ckpt_dir must be a directory.
     prefix: str: name prefix of checkpoint files.
     parallel: bool: whether to load seekable checkpoints in parallel, for speed.


Test patch: