Repository: google/flax

Issue:
GCS-related error during deletion of old checkpoints with save_checkpoint
### System information
- OS Platform and Distribution: ```TPU VM```.
- Flax, jax, jaxlib versions: ```flax==0.6.1, jax==0.3.21, jaxlib==0.3.20```
- Python version: ```3.8.10```
- GPU/TPU model and memory: ```TPU v3-8```
- CUDA version (if applicable): ```N/A```


### Problem you have encountered:

When running ```flax.training.checkpoints.save_checkpoint```, I am encountering an error when flax attempts to delete older checkpoints from Google Cloud Storage, complaining that they aren't directories. This seems to happen regardless of how flat or hierarchical the namespace for checkpoints is. 

The offending function is ```gfile.rmtree``` as called by ```flax.training.checkpoints._remove_invalid_ckpts```. 
Perhaps it should be ```gfile.remove```? 

### What you expected to happen:

It should delete the checkpoints without an error. 

### Logs, error messages, etc:

```
Traceback (most recent call last):
  File "scripts/launch.py", line 659, in <module>
    raise e
  File "scripts/launch.py", line 653, in <module>
    main()
  File "scripts/launch.py", line 618, in main
    train(
  File "scripts/launch.py", line 400, in train
    save_checkpoint(
  File "/home/lucaslingle/project123/src/project123/utils/persist_ops.py", line 30, in save_checkpoint
    checkpoints.save_checkpoint(
  File "/home/lucaslingle/.local/lib/python3.8/site-packages/flax/training/checkpoints.py", line 543, in save_checkpoint
    save_main_ckpt_task()
  File "/home/lucaslingle/.local/lib/python3.8/site-packages/flax/training/checkpoints.py", line 537, in save_main_ckpt_task
    return _save_main_ckpt_file(target, False, (ckpt_tmp_path, ckpt_path),
  File "/home/lucaslingle/.local/lib/python3.8/site-packages/flax/training/checkpoints.py", line 461, in _save_main_ckpt_file
    _save_commit(
  File "/home/lucaslingle/.local/lib/python3.8/site-packages/flax/training/checkpoints.py", line 420, in _save_commit
    _remove_invalid_ckpts(ckpt_path, base_path, keep, overwrite,
  File "/home/lucaslingle/.local/lib/python3.8/site-packages/flax/training/checkpoints.py", line 380, in _remove_invalid_ckpts
    gfile.rmtree(path)
  File "/home/lucaslingle/.local/lib/python3.8/site-packages/tensorflow/python/lib/io/file_io.py", line 674, in delete_recursively_v2
    _pywrap_file_io.DeleteRecursively(compat.path_to_bytes(path))
tensorflow.python.framework.errors_impl.NotFoundError: gs://bucket_name/model_name/checkpoints_10 doesn't exist or not a directory.
Traceback (most recent call last):
  File "scripts/launch.py", line 659, in <module>
    raise e
  File "scripts/launch.py", line 653, in <module>
    main()
  File "scripts/launch.py", line 618, in main
    train(
  File "scripts/launch.py", line 400, in train
    save_checkpoint(
  File "/home/lucaslingle/project123/src/project123/utils/persist_ops.py", line 30, in save_checkpoint
    checkpoints.save_checkpoint(
  File "/home/lucaslingle/.local/lib/python3.8/site-packages/flax/training/checkpoints.py", line 543, in save_checkpoint
    save_main_ckpt_task()
  File "/home/lucaslingle/.local/lib/python3.8/site-packages/flax/training/checkpoints.py", line 537, in save_main_ckpt_task
    return _save_main_ckpt_file(target, False, (ckpt_tmp_path, ckpt_path),
  File "/home/lucaslingle/.local/lib/python3.8/site-packages/flax/training/checkpoints.py", line 461, in _save_main_ckpt_file
    _save_commit(
  File "/home/lucaslingle/.local/lib/python3.8/site-packages/flax/training/checkpoints.py", line 420, in _save_commit
    _remove_invalid_ckpts(ckpt_path, base_path, keep, overwrite,
  File "/home/lucaslingle/.local/lib/python3.8/site-packages/flax/training/checkpoints.py", line 380, in _remove_invalid_ckpts
    gfile.rmtree(path)
  File "/home/lucaslingle/.local/lib/python3.8/site-packages/tensorflow/python/lib/io/file_io.py", line 674, in delete_recursively_v2
    _pywrap_file_io.DeleteRecursively(compat.path_to_bytes(path))
tensorflow.python.framework.errors_impl.NotFoundError: gs://bucket_name/model_name/checkpoints_10 doesn't exist or not a directory.
```
Note that the original bucket_name and model_name have been redacted, and that the checkpoints in question were uploaded to (and are still present in) the Google cloud storage bucket.

### Steps to reproduce:

```
from flax.training import checkpoints
import jax.numpy as jnp

SAVE_DIR = "gs://BUCKET_NAME"

checkpoints.save_checkpoint(
    ckpt_dir=SAVE_DIR,
    target={"params": jnp.array(10)},  # it should work on any pytree
    step=0,
    keep=1,
)

checkpoints.save_checkpoint(
    ckpt_dir=SAVE_DIR,
    target={"params": 2 * jnp.array(10)},  # it should work on any pytree
    step=1,
    keep=1,
)
```



Discussion:

Resolving patch:
diff --git a/flax/linen/stochastic.py b/flax/linen/stochastic.py
--- a/flax/linen/stochastic.py
+++ b/flax/linen/stochastic.py
@@ -29,7 +29,7 @@ class Dropout(Module):
 
     Note: When using :meth:`Module.apply() <flax.linen.Module.apply>`, make sure
     to include an RNG seed named `'dropout'`. For example::
-    
+
       model.apply({'params': params}, inputs=inputs, train=True, rngs={'dropout': dropout_rng})`
 
     Attributes:
diff --git a/flax/struct.py b/flax/struct.py
--- a/flax/struct.py
+++ b/flax/struct.py
@@ -98,7 +98,7 @@ def create(cls, kernel):
   # check if already a flax dataclass
   if '_flax_dataclass' in clz.__dict__:
     return clz
-      
+
   data_clz = dataclasses.dataclass(frozen=True)(clz)
   meta_fields = []
   data_fields = []
diff --git a/flax/training/checkpoints.py b/flax/training/checkpoints.py
--- a/flax/training/checkpoints.py
+++ b/flax/training/checkpoints.py
@@ -353,7 +353,7 @@ def _remove_invalid_ckpts(ckpt_path: str, base_path: str, keep: int,
         # checkpoint folder and before deleting the main checkpoint.
         if gfile.exists(path + MP_ARRAY_POSTFIX):
           gfile.rmtree(path + MP_ARRAY_POSTFIX)
-      gfile.rmtree(path)
+      gfile.remove(path)
 
   # Remove old checkpoint files.
   last_kept = -float('inf')
@@ -374,7 +374,7 @@ def _remove_invalid_ckpts(ckpt_path: str, base_path: str, keep: int,
         # MPA might be removed already but the main ckpt is still there.
         if gfile.exists(path + MP_ARRAY_POSTFIX):
           gfile.rmtree(path + MP_ARRAY_POSTFIX)
-      gfile.rmtree(path)
+      gfile.remove(path)
 
 
 def _save_commit(ckpt_tmp_path: str, ckpt_path: str, base_path: str, keep: int,


Test patch:
diff --git a/tests/struct_test.py b/tests/struct_test.py
--- a/tests/struct_test.py
+++ b/tests/struct_test.py
@@ -67,7 +67,7 @@ def test_keypath_error(self):
       raise e('in_axes')
 
   def test_double_wrap_no_op(self):
-    
+
     class A:
       a: int