Repository navigation
Make orbax-export an optional dependency? #5352
Description
Dear flax team,
We noticed that the latest flax requires orbax-export as a dependency (see PR #5188). orbax-export seems to require with TensorFlow, which is why currently its conda-forge build fails: conda-forge/staged-recipes#32625 (@matthewfeickert can comment more on this).
To make flax more lightweight on 'default' installation we were wondering what you think about making orbax-export an optional dependency to not make TensorFlow be a requirement for flax through orbax-export.
What do you think about something like pip install flax[export] ?
Best, Peter
Reacted by Matthew Feickert
Activity
vfdev-5 commentedon Mar 18, 2026
@pfackeldey thanks for the report! I agree and think that orbax-export should either testing dependency or installed manually in the notebook. @samanklesaria was it intentional to put orbax-export as hard dependency?
matthewfeickert commentedon Mar 18, 2026
(Currently at a conference, so replying quickly now and then more later) c.f. conda-forge/flax-feedstock#61 (comment) for some context on things on the conda-forge distribution side.
The real issue is that
orbax-exportis not programmatically transparent about its runtime requirements (import orbax.exportimportstensorflowand so atpip install orbax-exporttimeimport orbax.exportis broken). This is both not good from a packaging standards view and I think fair to say that most people using JAX-based tooling don't expect (or want) TensorFlow to be a hard requirement.samanklesaria commentedon Mar 18, 2026
It was indeed not intended to make orbax-export a hard dependency. Apologies- will fix promptly!
matthewfeickert commentedon Mar 18, 2026
Thanks. This is useful to know as for conda-forge builds we'll just patch that out of the
pyproject.tomlfor all the affected releases (v0.12.3tov0.12.5).matthewfeickert commentedon Mar 26, 2026
Note that PR #5353 didn't land until after https://github.com/google/flax/releases/tag/v0.12.6, and so
flax/pyproject.toml
Line 25 in 1572c84
is still a requirement in the latest
flaxrelease. This will get patched out in the conda-forge distribution (conda-forge/flax-feedstock#67), butflaxwithouttensorflowas a requirement won't be available on PyPI untilv0.12.7+.