xax 0.2.4__py3-none-any.whl → 0.2.5__py3-none-any.whl

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/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/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.dist-info → xax-0.2.5.dist-info}/METADATA

@@ -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.dist-info → xax-0.2.5.dist-info}/RECORD

@@ -1,14 +1,14 @@
-xax/__init__.py,sha256=bMKUtRtmVnHshkD4Ylw7ymzIPpcasJAoXnBsIdoSEng,14225
+xax/__init__.py,sha256=X_QqDNJir1wdsfRY1CU1F4mdCQMlMZnyqPtY8MM1ODU,14225
 xax/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/requirements-dev.txt,sha256=qkscNkFzWd1S5fump-AKH53rR65v2x5FmboFdy_kKvs,128
 xax/requirements.txt,sha256=6qY-84e-sTmlfJNrSjwONQKqzAn5h8G_oGIhnhmfSr4,302
 xax/core/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/core/conf.py,sha256=Wuo5WLRWuRTgb8eaihvnG_NZskTu0-P3JkIcl_hKINM,5124
-xax/core/state.py,sha256=bJONQ0wXgbgo1jjSqV3JtqG5tdMlli93Nax_ftZ2D0w,4552
+xax/core/state.py,sha256=yO25lMoLCUTJlHyLzQxlDbsHC_GZ3HkrKAq5huA7AkU,4552
 xax/nn/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/nn/embeddings.py,sha256=bQGxBFxkLwi2MQLkRfGaHPH5P_KKB21HdI7VNWTKIOQ,11847
 xax/nn/equinox.py,sha256=5fdOKRXqAVZPsV-aEez3i1wamr_oBYnG74GP1jEthjM,4843
-xax/nn/export.py,sha256=7Yemw3T33QGEP8RkmTkpu6tRVOhut2RUJmttNFfCgFw,5537
+xax/nn/export.py,sha256=pRfM2B4hB2EvljysC6AjtgB_7Cn7JtaP3dhYU2stZtY,5545
 xax/nn/functions.py,sha256=CI_OmspaQwN9nl4hwefIU3_I7m6gBZwJ9aGK1JGUgr0,2713
 xax/nn/geom.py,sha256=rImNlkHWeoNcY7f84nknizJ6uzsrMhbAtKeb2xAWxNY,6215
 xax/nn/losses.py,sha256=Q_NVnm5n4UPBvp5nI_1aUptfXnqFYoUeFwySiyvopHg,272
@@ -58,8 +58,8 @@ xax/utils/data/collate.py,sha256=Rd9vMomr_S_zCa_Hi4dO-8ntzAfVwndIUtuXFA3iNcc,706
 xax/utils/types/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/utils/types/frozen_dict.py,sha256=ZCMGfSfr2_b2qZbq9ywPD0zej5tpVSId2JftXpwfB5k,4686
 xax/utils/types/hashable_array.py,sha256=l5iIcFmkYzfGeaZmcSoeFkthFASqM8xJYK3AXhZQYwc,992
-xax-0.2.4.dist-info/licenses/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
-xax-0.2.4.dist-info/METADATA,sha256=9hMsPCoszpjVN0rLDMlT20aYqmQwnHvl9T1V_0akl0U,1882
-xax-0.2.4.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-xax-0.2.4.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
-xax-0.2.4.dist-info/RECORD,,
+xax-0.2.5.dist-info/licenses/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
+xax-0.2.5.dist-info/METADATA,sha256=4RBxZF_P0cg-a6QUNS9urvzc4BGGfoedqMrnP0L6Ksk,1879
+xax-0.2.5.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
+xax-0.2.5.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
+xax-0.2.5.dist-info/RECORD,,