xax 0.1.5__py3-none-any.whl → 0.1.6__py3-none-any.whl

xax/__init__.py

@@ -12,7 +12,7 @@ and running the update script:
     python -m scripts.update_api --inplace
 """
 
-__version__ = "0.1.5"
+__version__ = "0.1.6"
 
 # This list shouldn't be modified by hand; instead, run the update script.
 __all__ = [
@@ -97,6 +97,8 @@ __all__ = [
     "save_config",
     "stage_environment",
     "to_markdown_table",
+    "HashableArray",
+    "hashable_array",
     "jit",
     "save_jaxpr_dot",
     "ColoredFormatter",
@@ -251,6 +253,8 @@ NAME_MAP: dict[str, str] = {
     "save_config": "utils.experiments",
     "stage_environment": "utils.experiments",
     "to_markdown_table": "utils.experiments",
+    "HashableArray": "utils.jax",
+    "hashable_array": "utils.jax",
     "jit": "utils.jax",
     "save_jaxpr_dot": "utils.jaxpr",
     "ColoredFormatter": "utils.logging",
@@ -404,7 +408,7 @@ if IMPORT_ALL or TYPE_CHECKING:
         stage_environment,
         to_markdown_table,
     )
-    from xax.utils.jax import jit
+    from xax.utils.jax import HashableArray, hashable_array, jit
     from xax.utils.jaxpr import save_jaxpr_dot
     from xax.utils.logging import (
         LOG_ERROR_SUMMARY,

xax/utils/jax.py

@@ -138,3 +138,25 @@ def jit(
         return wrapped
 
     return decorator
+
+
+class HashableArray:
+    def __init__(self, array: np.ndarray | jnp.ndarray) -> None:
+        if not isinstance(array, (np.ndarray, jnp.ndarray)):
+            raise ValueError(f"Expected np.ndarray or jnp.ndarray, got {type(array)}")
+        self.array = array
+        self._hash: int | None = None
+
+    def __hash__(self) -> int:
+        if self._hash is None:
+            self._hash = hash(self.array.tobytes())
+        return self._hash
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, HashableArray):
+            return False
+        return bool(jnp.array_equal(self.array, other.array))
+
+
+def hashable_array(array: np.ndarray | jnp.ndarray) -> HashableArray:
+    return HashableArray(array)

{xax-0.1.5.dist-info → xax-0.1.6.dist-info}/METADATA

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

{xax-0.1.5.dist-info → xax-0.1.6.dist-info}/RECORD

@@ -1,4 +1,4 @@
-xax/__init__.py,sha256=Ig2JALbHQOtFHr57jpmPm16L0aYkNTNZ-Upz9jxIyrc,13361
+xax/__init__.py,sha256=rjqydWhxQVUAj3lXgFpzj4iLFOdDJGHArfAH7_QSkhk,13504
 xax/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/requirements-dev.txt,sha256=qkscNkFzWd1S5fump-AKH53rR65v2x5FmboFdy_kKvs,128
 xax/requirements.txt,sha256=9LAEZ5c5gqRSARRVA6xJsVTa4MebPZuC4yOkkwkZJFw,297
@@ -43,7 +43,7 @@ xax/task/mixins/train.py,sha256=vsH_QpyrThlh9AzWnyvDJv58Y8U_516oi8gmMq_0iMg,2233
 xax/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/utils/debugging.py,sha256=9WlCrEqbq-SVXPEM4rhsLYERH97XNX7XSYLSI3sgKGk,1619
 xax/utils/experiments.py,sha256=5CUja1H_cx4dnVqTGQekOpIhqISwHtAgLxZ34GV7cwM,29229
-xax/utils/jax.py,sha256=tC0NNelbrSTzwNGluiwLGKtoHhVpgdzrv-xherB3VtY,4752
+xax/utils/jax.py,sha256=eObvWt2DraCs2IMDZSdQ0rRk8tA3P5XBlF_UeVq7Aro,5480
 xax/utils/jaxpr.py,sha256=S80nyEkv188RInzq3kCAdkQCU-bf6s0oPTrCE_LjkRs,2298
 xax/utils/logging.py,sha256=GAhTne2rdB4Fa1lzk06DMO15U8MTejn6XTClShC-ZtU,6622
 xax/utils/numpy.py,sha256=_jOXVi-d2AtJnRftPkRK5MDMzsU8slgw-Jjv4GRm6ns,1197
@@ -53,8 +53,8 @@ xax/utils/tensorboard.py,sha256=21czW8WC2SAmwEhz6RLJc_q5HFvNKM4iR1ZycSO5qPE,1705
 xax/utils/text.py,sha256=zo1sAoZe59GkpcpaHBVOQ0OekSMGXvOAyNa3lOJozCY,10628
 xax/utils/data/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/utils/data/collate.py,sha256=Rd9vMomr_S_zCa_Hi4dO-8ntzAfVwndIUtuXFA3iNcc,7066
-xax-0.1.5.dist-info/licenses/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
-xax-0.1.5.dist-info/METADATA,sha256=dkvPH-GXRErGzzR0VGsiRisNYBzDVf97nEGG0F2-HjY,1877
-xax-0.1.5.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-xax-0.1.5.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
-xax-0.1.5.dist-info/RECORD,,
+xax-0.1.6.dist-info/licenses/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
+xax-0.1.6.dist-info/METADATA,sha256=vKxhuOt02ALjFV9fAt-rPVTwvqX4uNr_shL1DGEotA4,1877
+xax-0.1.6.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
+xax-0.1.6.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
+xax-0.1.6.dist-info/RECORD,,