xax 0.3.9__py3-none-any.whl → 0.3.10__py3-none-any.whl

xax/__init__.py

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

xax/utils/pytree.py

@@ -1,6 +1,6 @@
 """Utils for accessing, modifying, and otherwise manipulating pytrees."""
 
-from typing import TypeVar
+from typing import Mapping, Sequence, TypeVar
 
 import chex
 import equinox as eqx
@@ -258,11 +258,18 @@ def tuple_insert(t: tuple[T, ...], index: int, value: T) -> tuple[T, ...]:
 def get_pytree_mapping(pytree: PyTree) -> dict[str, Array]:
     leaves: dict[str, Array] = {}
 
+    def _get_str(thing: PyTree) -> str:
+        if isinstance(thing, str):
+            return thing
+        if isinstance(thing, Sequence):
+            return "/".join(_get_str(x) for x in thing)
+        if isinstance(thing, Mapping):
+            return "/".join(f"{_get_str(k)}:{_get_str(v)}" for k, v in thing.items())
+        return str(thing)
+
     def _get_leaf(path: tuple, x: PyTree) -> None:
         if isinstance(x, jnp.ndarray):
-            # Convert path tuple to string, e.g. (1, 'a', 2) -> '1/a/2'
-            path_str = "/".join(str(p) for p in path)
-            leaves[path_str] = x
+            leaves[_get_str(path)] = x
 
     jax.tree.map_with_path(_get_leaf, pytree)
     return leaves

{xax-0.3.9.dist-info → xax-0.3.10.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: xax
-Version: 0.3.9
+Version: 0.3.10
 Summary: A library for fast Jax experimentation
 Home-page: https://github.com/kscalelabs/xax
 Author: Benjamin Bolte

{xax-0.3.9.dist-info → xax-0.3.10.dist-info}/RECORD

@@ -1,4 +1,4 @@
-xax/__init__.py,sha256=PLNyZ5fOm0f2JduTMauNH2jqxN4g0GAZ9MSNzEsAQS4,16665
+xax/__init__.py,sha256=lSwyrPTof_BZ-pyPNhNICJnCZMN9i2sJ-Ii3S_vY_28,16666
 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
@@ -52,7 +52,7 @@ xax/utils/jaxpr.py,sha256=H7pWl48ROXIB1-ZPWYfOn-ou3EBMxYWIwc_A0reJQoo,2333
 xax/utils/logging.py,sha256=Kkyma_LJXqrN2HTQ214gRP_9ih3_bKk115MWC60lQWM,6656
 xax/utils/numpy.py,sha256=_jOXVi-d2AtJnRftPkRK5MDMzsU8slgw-Jjv4GRm6ns,1197
 xax/utils/profile.py,sha256=-aFdWpgYFvBsBZXSLL4zXrFe3zzsDqzmx4q5f2WOtpQ,1628
-xax/utils/pytree.py,sha256=cLZRSd5xc-DqcbRfWnBy87pAiUU5fT8U4CHoLi_i_v4,9642
+xax/utils/pytree.py,sha256=w8Ab2LmJdQ8e1FxKF0xWWaOak09Mhu44ZcOeUR6uGFA,9889
 xax/utils/tensorboard.py,sha256=P0oIFvX2Qts1H4lkpizhRIpQdD0MNppVMeut0Z94yCs,19878
 xax/utils/text.py,sha256=xS02aSzdywl3KIaNSpKWcxdd37oYlUJtu9wIjkc1wVc,10654
 xax/utils/data/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -60,9 +60,9 @@ 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=ebtHENhyUzSjyJTlbMaLtcckQIJ7EtgJiok_40TJZpo,4689
 xax/utils/types/hashable_array.py,sha256=l5iIcFmkYzfGeaZmcSoeFkthFASqM8xJYK3AXhZQYwc,992
-xax-0.3.9.dist-info/licenses/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
-xax-0.3.9.dist-info/METADATA,sha256=UKUqiAADutUzjtq4WpChJJzWDcfKBsrbWMmZu0LUqQk,1246
-xax-0.3.9.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-xax-0.3.9.dist-info/entry_points.txt,sha256=uRC6rx5ce0bf-FblJaZSBMxxKFfMyoWTf8OWbBmLSe8,61
-xax-0.3.9.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
-xax-0.3.9.dist-info/RECORD,,
+xax-0.3.10.dist-info/licenses/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
+xax-0.3.10.dist-info/METADATA,sha256=oQMGYjsfYxMmw0A60qE15yda_G-0YG5RNl17tboR1f0,1247
+xax-0.3.10.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+xax-0.3.10.dist-info/entry_points.txt,sha256=uRC6rx5ce0bf-FblJaZSBMxxKFfMyoWTf8OWbBmLSe8,61
+xax-0.3.10.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
+xax-0.3.10.dist-info/RECORD,,