TensorFlow is not listed as a `setup.py` dependency but is required

[samuela]

Problem you have encountered:

tensorflow is not listed as a dependency in setup.py here: https://github.com/google/flax/blob/35d5d8a4ff466f57bf6af4d0f8c147d472795743/setup.py#L27-L33 but is required by flax here: https://github.com/google/flax/blob/aad0be1e9b20e3a571c9a6d7814bda7a9951ba5c/flax/training/checkpoints.py#L31

What you expected to happen:

tensorflow to be listed as a requirement in setup.py. Or, even better, for there to be no dependence on tensorflow at all since it seems to be only used for tensorflow.io.gfile.

Logs, error messages, etc:

n/a

Steps to reproduce:

n/a

[PhilipVinc]

Please don't make flax depend on tensorflow. This can easily be an extra_requires dependency.

[samuela]

Agreed, but I think optimal solution here is to not use tensorflow.gfile at all.

[jheek]

I think what we should do is fall back to using python file API if tensorflow is not installed.

[AdityaKane2001]

I think what we should do is fall back to using python file API if tensorflow is not installed.

IMHO this should be the default method. One should need tf if and only if GCS is involved.

[levskaya]

#2073 is a WIP first draft at removing the TF-dependence for basic IO operations. We can't shim something that will perform as well as the native gfile io calls, but we can break the dependency so things at least work in most simple use-cases without a TF install.

We would still have an implicit dependence on TF for our metrics/tensorboard.py tensorboard utility though. The tensorboard people finally released a version of tensorboard that can be installed without tensorflow, though they've written only a little about what's supported in this mode of operation. Still, perhaps we can just deprecate our own summary writer util at this point and rely on them for tensorboard metric export.

[marcvanzee]

@IvyZX maybe you could chime in here and tell us what the plan is for depending on TF for checkpointing? It seems we currently only depend on it for tf_errors, is this temporary?

[IvyZX]

PR https://github.com/google/flax/pull/2697 is moving tf_errors inside the file system shim io.py and dropping TF from the require import list of flax.training.checkpoints.

Now the metrics/tensorboard.py should be the only place where TF is imported. Tensorboard is a much lighter dependency and seems to support all the tf.summary use cases inside metrics/tensorboard.py, so in the future perhaps user could directly use tensorboard for metric writers.