xax 0.1.3__py3-none-any.whl → 0.1.4__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.3"
+__version__ = "0.1.4"
 
 # This list shouldn't be modified by hand; instead, run the update script.
 __all__ = [

xax/task/base.py

@@ -153,6 +153,8 @@ class BaseTask(Generic[Config]):
             for base in cls.__orig_bases__:
                 if hasattr(base, "__args__"):
                     for arg in base.__args__:
+                        if isinstance(arg, TypeVar) and arg.__bound__ is not None:
+                            arg = arg.__bound__
                         if issubclass(arg, BaseConfig):
                             return arg
 

xax/task/mixins/compile.py

@@ -5,6 +5,7 @@ behavior during initialization and training.
 """
 
 import logging
+import sys
 from dataclasses import dataclass
 from pathlib import Path
 from typing import Generic, TypeVar
@@ -17,6 +18,16 @@ from xax.task.base import BaseConfig, BaseTask
 logger = logging.getLogger(__name__)
 
 
+def get_cache_dir() -> str | None:
+    # By default, only cache on MacOS, since Jax caching on Linux is very
+    # prone to NaNs.
+    match sys.platform:
+        case "darwin":
+            return str((Path.home() / ".cache" / "jax" / "jaxcache").resolve())
+        case _:
+            return None
+
+
 @jax.tree_util.register_dataclass
 @dataclass
 class CompileOptions:
@@ -42,7 +53,8 @@ class CompileOptions:
 
     # JAX cache options
     cache_dir: str | None = field(
-        value=lambda: str((Path.home() / ".cache" / "jax" / "jaxcache").resolve()),
+        # Only cache by default on MacOS systems.
+        value=get_cache_dir,
         help="Directory for JAX compilation cache. If None, caching is disabled",
     )
     cache_min_size_bytes: int = field(
@@ -54,7 +66,7 @@ class CompileOptions:
         help="Minimum compilation time in seconds for cache entries. 0 means no minimum",
     )
     cache_enable_xla: str = field(
-        value="all",
+        value="none",
         help="Which XLA caches to enable",
     )
 

{xax-0.1.3.dist-info → xax-0.1.4.dist-info}/METADATA

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

{xax-0.1.3.dist-info → xax-0.1.4.dist-info}/RECORD

@@ -1,4 +1,4 @@
-xax/__init__.py,sha256=n7EXl0pwEPzlw2DjS-3ePgx0VoQnMDnHLVc5exkHGcM,13361
+xax/__init__.py,sha256=Ib-tn4sRqSIiFjoQceuUzdwRc4i-VW9P5bBLq4wDbYc,13361
 xax/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/requirements-dev.txt,sha256=qkscNkFzWd1S5fump-AKH53rR65v2x5FmboFdy_kKvs,128
 xax/requirements.txt,sha256=9LAEZ5c5gqRSARRVA6xJsVTa4MebPZuC4yOkkwkZJFw,297
@@ -14,7 +14,7 @@ xax/nn/geom.py,sha256=eK7I8fUHBc3FT7zpm5Yf__bXFQ4LtX6sa17-DxojLTo,3202
 xax/nn/norm.py,sha256=cDmYf5CtyzmuCiWdSP5nr8nZKQOmaZueDQXMPnThg6c,548
 xax/nn/parallel.py,sha256=fnTiT7MsG7eQrJvqwjIz2Ifo3P27TuxIJzmpGYSa_dQ,4608
 xax/task/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-xax/task/base.py,sha256=MlH5dTKAiMzFRI5fmXCvL1k8ELbalWMBICeVxmW6k2U,7479
+xax/task/base.py,sha256=4fUjrG-llQpeESQuaQbww4M6WR6djjTK89fY20UV9zU,7610
 xax/task/logger.py,sha256=1SZjVC6UCtZUoMPcpp3ckotL324QDeYDvHVhf5MHVqg,36271
 xax/task/script.py,sha256=zt36Sobdoer86gXHqc4sMAW7bqZRVl6IEExuQZH2USk,926
 xax/task/task.py,sha256=UHMpnv__gqMcfbC_L-Hhk-DCnUYlFVsgbNf-v8o8B7U,1424
@@ -31,7 +31,7 @@ xax/task/loggers/tensorboard.py,sha256=kI8LvBuBBhPgkP8TeaTQb9SQ0FqaIodwQh2SuWDCn
 xax/task/mixins/__init__.py,sha256=D3oU31rB9FeOr9MPLleLt5JFbftUr4sBTwgnwQdc2qA,809
 xax/task/mixins/artifacts.py,sha256=2ezmZGzPGe3nhsd9KRkeHWWXdbT9m7drzimIfw6v1XY,2892
 xax/task/mixins/checkpointing.py,sha256=a6tVyISsDIz68rrhb1rAh3rjQlqkDVJCmSBmETQrnRM,8480
-xax/task/mixins/compile.py,sha256=FRsxwLnZjjxpeWJ7Bx_d8XUY50oDoGidgpeRt4ejeQk,3377
+xax/task/mixins/compile.py,sha256=9pVJEUvizu6-6tq0HaMtHGNSi9Yk_mxNyqBFimcfwL0,3683
 xax/task/mixins/cpu_stats.py,sha256=C_t71UTrv4LwQzhO5iubsfomj4jYa9bzpE4zBcHdoHM,9211
 xax/task/mixins/data_loader.py,sha256=WjMWk9uACfBMMClLMcLPkE0WNIvlCZnmqyyqLqJpjX0,6545
 xax/task/mixins/gpu_stats.py,sha256=IGPBro9xzSivwD43zM18lWcuei7IhA8LilxSPHqNl4I,8747
@@ -53,8 +53,8 @@ xax/utils/tensorboard.py,sha256=_S70dS69pduiD05viHAGgYGsaBry1QL2ej6ZwUIXPOE,1617
 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.3.dist-info/licenses/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
-xax-0.1.3.dist-info/METADATA,sha256=m3AyjlRD9C-O2Tp5zH5i5TbEL7bZooeIpypUCYuYPtQ,1877
-xax-0.1.3.dist-info/WHEEL,sha256=DK49LOLCYiurdXXOXwGJm6U4DkHkg4lcxjhqwRa0CP4,91
-xax-0.1.3.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
-xax-0.1.3.dist-info/RECORD,,
+xax-0.1.4.dist-info/licenses/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
+xax-0.1.4.dist-info/METADATA,sha256=VMyCF4gPJh5ww_VzF0HgfInf_bfD40ahwxpHkxEMGoI,1877
+xax-0.1.4.dist-info/WHEEL,sha256=DK49LOLCYiurdXXOXwGJm6U4DkHkg4lcxjhqwRa0CP4,91
+xax-0.1.4.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
+xax-0.1.4.dist-info/RECORD,,