up

Author: sayakpaul

setup.py

@@ -251,7 +251,7 @@ def run(self):
     extras["flax"] = deps_list("jax", "jaxlib", "flax")
 
 
-extras["dev"] = extras["quality"] + extras["test"] + extras["training"] + extras["docs"] + extras["torch"]
+extras["dev"] = extras["quality"] + extras["test"] + extras["training"] + extras["docs"] + extras["torch"] + extras["flax"]
 
 install_requires = [
     deps["importlib_metadata"],