Repository: google/flax

Issue:
Tensorflow Required Dependancy For Saving Checkpoint?
Hi,

I tried following the example in https://github.com/google/flax/discussions/1876#discussioncomment-2143866 and ran into an ModuleNotFoundError. I was able to fix it by install tensorflow, but was wondering if tensorflow should be a required dependancy.

### System information
- OS Platform and Distribution (e.g., Linux Ubuntu 16.04): 
- Flax, jax, jaxlib versions (obtain with `pip show flax jax jaxlib`: 
```
Name: flax Version: 0.6.0
Name: jax Version: 0.3.17
Name: jaxlib Version: 0.3.15
```
- Python version: python=3.8
- GPU/TPU model and memory:  CPU


### Problem you have encountered:
ModuleNotFoundError: No module named tensorflow

### What you expected to happen:
checkpoints should be able to imported without throwing any errors

### Logs, error messages, etc:
```
File "test.py", line 5, in <module>
    from flax.training import train_state, checkpoints
  File "/home/user/anaconda3/envs/jax-dep/lib/python3.8/site-packages/flax/training/checkpoints.py", line 36, in <module>
    from tensorflow.io import gfile  # pytype: disable=import-error
ModuleNotFoundError: No module named 'tensorflow'
``` 
### Steps to reproduce:
```
import jax.numpy as jnp
import jax
import flax
import flax.linen as nn
from flax.training import train_state, checkpoints
import optax
import numpy as np


class Net(nn.Module):
        features: int

        @nn.compact
        def __call__(self, x):
            x = nn.Dense(self.features)(x)
            x = nn.Dense(self.features)(x)
            return x


model = Net(features=2)
params = model.init(jax.random.PRNGKey(0), jnp.ones((1, 2)))

tx = optax.adam(learning_rate=0.0001)
state = train_state.TrainState.create(apply_fn=model.apply, params=params, tx=tx)


CKPT_DIR = 'ckpts'
checkpoints.save_checkpoint(ckpt_dir=CKPT_DIR, target=state, step=1)
restored_state = checkpoints.restore_checkpoint(ckpt_dir=CKPT_DIR, target=state)

assert jax.tree_util.tree_all(jax.tree_util.tree_map(lambda x, y: (x == y).all(), state.params, restored_state.params))
```


Discussion:
There is work in progress to remove TensorFlow as a dependency in #2073 and #2444.
This should be solved in the future.
On Sep 16, d9e6e44b545b7805ab827be499b2afa689f7e673 (v0.6.1) added another tensorflow dependency. Just commenting that neither #2073 and #2444 (although not complete yet) is aware of this change, so we look forward to have the tensorflow dependency removed.

I can also provide a help by writing some PRs; it would be great to systemically prevent such accidental use of `tensorflo` (in setup.py).

/cc @IvyZX the owner of the line.
@sudo-michael @wookayin Check out the new Checkpointing guide by @IvyZX 

https://github.com/google/flax/blob/main/docs/guides/use_checkpointing.ipynb
UPDATES: #2586 has been merged.

Resolving patch:
diff --git a/flax/io.py b/flax/io.py
--- a/flax/io.py
+++ b/flax/io.py
@@ -47,6 +47,16 @@ class BackendMode(Enum):
   io_mode = BackendMode.DEFAULT
 
 
+# Constants and Exceptions
+
+
+if io_mode == BackendMode.TF:
+  from tensorflow import errors as tf_errors  # type: ignore
+  NotFoundError = tf_errors.NotFoundError
+else:
+  NotFoundError = FileNotFoundError
+
+
 # Overrides for testing.
 
 
@@ -176,4 +186,4 @@ def rmtree(path):
   elif io_mode == BackendMode.TF:
     return gfile.rmtree(path)
   else:
-    raise ValueError("Unknown IO Backend Mode.")
\ No newline at end of file
+    raise ValueError("Unknown IO Backend Mode.")
diff --git a/flax/training/checkpoints.py b/flax/training/checkpoints.py
--- a/flax/training/checkpoints.py
+++ b/flax/training/checkpoints.py
@@ -39,7 +39,6 @@
 from jax.experimental.global_device_array import GlobalDeviceArray
 from jax.experimental.multihost_utils import sync_global_devices
 import orbax.checkpoint as orbax
-from tensorflow import errors as tf_errors
 
 
 _IMPORT_GDAM_SUCCESSFUL = False
@@ -107,7 +106,7 @@ def _checkpoint_path_step(path: str) -> Optional[float]:
 def _allowempty_listdir(path: str):
   try:
     return io.listdir(path)
-  except tf_errors.NotFoundError:
+  except io.NotFoundError:
     return []
 
 def _safe_remove(path: str):


Test patch: