xax 0.2.4__tar.gz → 0.2.5__tar.gz

{xax-0.2.4/xax.egg-info → xax-0.2.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: xax
-Version: 0.2.4
+Version: 0.2.5
 Summary: A library for fast Jax experimentation
 Home-page: https://github.com/kscalelabs/xax
 Author: Benjamin Bolte
@@ -31,11 +31,10 @@ Requires-Dist: pytest; extra == "dev"
 Requires-Dist: types-pillow; extra == "dev"
 Requires-Dist: types-psutil; extra == "dev"
 Requires-Dist: types-requests; extra == "dev"
-Provides-Extra: export
-Requires-Dist: orbax-export; extra == "export"
-Requires-Dist: tensorflow; extra == "export"
-Provides-Extra: flax
-Requires-Dist: flax; extra == "flax"
+Provides-Extra: exportable
+Requires-Dist: flax; extra == "exportable"
+Requires-Dist: orbax-export; extra == "exportable"
+Requires-Dist: tensorflow; extra == "exportable"
 Provides-Extra: all
 Requires-Dist: black; extra == "all"
 Requires-Dist: darglint; extra == "all"
@@ -45,9 +44,9 @@ Requires-Dist: pytest; extra == "all"
 Requires-Dist: types-pillow; extra == "all"
 Requires-Dist: types-psutil; extra == "all"
 Requires-Dist: types-requests; extra == "all"
+Requires-Dist: flax; extra == "all"
 Requires-Dist: orbax-export; extra == "all"
 Requires-Dist: tensorflow; extra == "all"
-Requires-Dist: flax; extra == "all"
 Dynamic: author
 Dynamic: description
 Dynamic: description-content-type

{xax-0.2.4 → xax-0.2.5}/setup.py

@@ -15,14 +15,11 @@ with open("xax/requirements-dev.txt", "r", encoding="utf-8") as f:
     requirements_dev: list[str] = f.read().splitlines()
 
 requirements_export: list[str] = [
+    "flax",
     "orbax-export",
     "tensorflow",
 ]
 
-requirements_flax: list[str] = [
-    "flax",
-]
-
 with open("xax/__init__.py", "r", encoding="utf-8") as fh:
     version_re = re.search(r"^__version__ = \"([^\"]*)\"", fh.read(), re.MULTILINE)
 assert version_re is not None, "Could not find version in xax/__init__.py"
@@ -42,9 +39,8 @@ setup(
     tests_require=requirements_dev,
     extras_require={
         "dev": requirements_dev,
-        "export": requirements_export,
-        "flax": requirements_flax,
-        "all": requirements_dev + requirements_export + requirements_flax,
+        "exportable": requirements_export,
+        "all": requirements_dev + requirements_export,
     },
     package_data={
         "xax": [

{xax-0.2.4 → xax-0.2.5}/xax/__init__.py

@@ -12,7 +12,7 @@ and running the update script:
     python -m scripts.update_api --inplace
 """
 
-__version__ = "0.2.4"
+__version__ = "0.2.5"
 
 # This list shouldn't be modified by hand; instead, run the update script.
 __all__ = [

{xax-0.2.4 → xax-0.2.5}/xax/core/state.py

@@ -89,9 +89,9 @@ class State:
             int32_arr = int32_arr.at[1].set(kwargs["num_valid_steps"])
 
         if "phase" in kwargs:
-            int32_arr = int32_arr.at[3].set(_phase_to_int(kwargs["phase"]))
+            int32_arr = int32_arr.at[2].set(_phase_to_int(kwargs["phase"]))
         if "_phase" in kwargs:
-            int32_arr = int32_arr.at[3].set(kwargs["_phase"])
+            int32_arr = int32_arr.at[2].set(kwargs["_phase"])
 
         if "num_samples" in kwargs:
             float32_arr = float32_arr.at[0].set(kwargs["num_samples"])

{xax-0.2.4 → xax-0.2.5}/xax/nn/export.py

@@ -14,8 +14,8 @@ try:
     from orbax.export import ExportManager, JaxModule, ServingConfig
 except ImportError as e:
     raise ImportError(
-        "In order to export models, please install Xax with export dependencies, "
-        "using 'xax[export]` to install the required dependencies."
+        "In order to export models, please install Xax with exportable dependencies, "
+        "using 'xax[exportable]` to install the required dependencies."
     ) from e
 
 logger = logging.getLogger(__name__)

{xax-0.2.4 → xax-0.2.5/xax.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: xax
-Version: 0.2.4
+Version: 0.2.5
 Summary: A library for fast Jax experimentation
 Home-page: https://github.com/kscalelabs/xax
 Author: Benjamin Bolte
@@ -31,11 +31,10 @@ Requires-Dist: pytest; extra == "dev"
 Requires-Dist: types-pillow; extra == "dev"
 Requires-Dist: types-psutil; extra == "dev"
 Requires-Dist: types-requests; extra == "dev"
-Provides-Extra: export
-Requires-Dist: orbax-export; extra == "export"
-Requires-Dist: tensorflow; extra == "export"
-Provides-Extra: flax
-Requires-Dist: flax; extra == "flax"
+Provides-Extra: exportable
+Requires-Dist: flax; extra == "exportable"
+Requires-Dist: orbax-export; extra == "exportable"
+Requires-Dist: tensorflow; extra == "exportable"
 Provides-Extra: all
 Requires-Dist: black; extra == "all"
 Requires-Dist: darglint; extra == "all"
@@ -45,9 +44,9 @@ Requires-Dist: pytest; extra == "all"
 Requires-Dist: types-pillow; extra == "all"
 Requires-Dist: types-psutil; extra == "all"
 Requires-Dist: types-requests; extra == "all"
+Requires-Dist: flax; extra == "all"
 Requires-Dist: orbax-export; extra == "all"
 Requires-Dist: tensorflow; extra == "all"
-Requires-Dist: flax; extra == "all"
 Dynamic: author
 Dynamic: description
 Dynamic: description-content-type

{xax-0.2.4 → xax-0.2.5}/xax.egg-info/requires.txt

@@ -23,9 +23,9 @@ pytest
 types-pillow
 types-psutil
 types-requests
+flax
 orbax-export
 tensorflow
-flax
 
 [dev]
 black
@@ -37,9 +37,7 @@ types-pillow
 types-psutil
 types-requests
 
-[export]
+[exportable]
+flax
 orbax-export
 tensorflow
-
-[flax]
-flax