Introduce flax.io.NotFoundError to remove tensorflow dependency
This fixes up a small mistake in #2586, where tensorflow is still being
imported from `flax.training.checkpoints`. To avoid the tensorflow
dependency, an alias `flax.io.NotFoundError` is added which can be
either `tf_errors.NotFoundError` or `FileNotFoundError` depending
on the backend mode.
