xax 0.1.13__tar.gz → 0.1.14__tar.gz

{xax-0.1.13/xax.egg-info → xax-0.1.14}/PKG-INFO

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

{xax-0.1.13 → xax-0.1.14}/xax/__init__.py

@@ -12,7 +12,7 @@ and running the update script:
     python -m scripts.update_api --inplace
 """
 
-__version__ = "0.1.13"
+__version__ = "0.1.14"
 
 # This list shouldn't be modified by hand; instead, run the update script.
 __all__ = [
@@ -77,6 +77,7 @@ __all__ = [
     "collate_non_null",
     "breakpoint_if_nan",
     "get_named_leaves",
+    "log_if_nan",
     "BaseFileDownloader",
     "ContextTimer",
     "CumulativeTimer",
@@ -237,6 +238,7 @@ NAME_MAP: dict[str, str] = {
     "collate_non_null": "utils.data.collate",
     "breakpoint_if_nan": "utils.debugging",
     "get_named_leaves": "utils.debugging",
+    "log_if_nan": "utils.debugging",
     "BaseFileDownloader": "utils.experiments",
     "ContextTimer": "utils.experiments",
     "CumulativeTimer": "utils.experiments",
@@ -388,7 +390,7 @@ if IMPORT_ALL or TYPE_CHECKING:
     from xax.task.script import Script, ScriptConfig
     from xax.task.task import Config, Task
     from xax.utils.data.collate import CollateMode, collate, collate_non_null
-    from xax.utils.debugging import breakpoint_if_nan, get_named_leaves
+    from xax.utils.debugging import breakpoint_if_nan, get_named_leaves, log_if_nan
     from xax.utils.experiments import (
         BaseFileDownloader,
         ContextTimer,

{xax-0.1.13 → xax-0.1.14}/xax/task/base.py

@@ -82,9 +82,6 @@ class BaseTask(Generic[Config]):
     def on_after_checkpoint_save(self, ckpt_path: Path, state: State) -> State:
         return state
 
-    def on_before_checkpoint_load(self, ckpt_path: Path) -> None:
-        pass
-
     @functools.cached_property
     def task_class_name(self) -> str:
         return self.__class__.__name__

{xax-0.1.13 → xax-0.1.14}/xax/task/mixins/checkpointing.py

@@ -98,19 +98,39 @@ class CheckpointingMixin(ArtifactsMixin[Config], Generic[Config]):
     ) -> tuple[PyTree, State, DictConfig]: ...
 
     @overload
-    def load_checkpoint(self, path: Path, part: Literal["model"]) -> PyTree: ...
+    def load_checkpoint(
+        self,
+        path: Path,
+        part: Literal["model"],
+    ) -> PyTree: ...
 
     @overload
-    def load_checkpoint(self, path: Path, part: Literal["opt"]) -> optax.GradientTransformation: ...
+    def load_checkpoint(
+        self,
+        path: Path,
+        part: Literal["opt"],
+    ) -> optax.GradientTransformation: ...
 
     @overload
-    def load_checkpoint(self, path: Path, part: Literal["opt_state"]) -> optax.OptState: ...
+    def load_checkpoint(
+        self,
+        path: Path,
+        part: Literal["opt_state"],
+    ) -> optax.OptState: ...
 
     @overload
-    def load_checkpoint(self, path: Path, part: Literal["state"]) -> State: ...
+    def load_checkpoint(
+        self,
+        path: Path,
+        part: Literal["state"],
+    ) -> State: ...
 
     @overload
-    def load_checkpoint(self, path: Path, part: Literal["config"]) -> DictConfig: ...
+    def load_checkpoint(
+        self,
+        path: Path,
+        part: Literal["config"],
+    ) -> DictConfig: ...
 
     def load_checkpoint(
         self,
@@ -125,9 +145,6 @@ class CheckpointingMixin(ArtifactsMixin[Config], Generic[Config]):
         | State
         | DictConfig
     ):
-        # Calls the base callback.
-        self.on_before_checkpoint_load(path)
-
         with tarfile.open(path, "r:gz") as tar:
 
             def get_model() -> PyTree:

{xax-0.1.13 → xax-0.1.14}/xax/task/mixins/compile.py

@@ -32,6 +32,10 @@ def get_cache_dir() -> str | None:
 @dataclass
 class CompileOptions:
     # JAX compilation options
+    debug_nans: bool = field(
+        value=False,
+        help="If True, breaks on NaNs",
+    )
     disable_jit: bool = field(
         value=False,
         help="If True, disables JIT compilation",
@@ -89,6 +93,10 @@ class CompileMixin(BaseTask[Config], Generic[Config]):
         cc = self.config.compile
 
         # Set basic compilation flags
+        if cc.debug_nans:
+            logger.info("Enabling NaNs debugging")
+            jax.config.update("jax_debug_nans", True)
+
         if cc.disable_jit:
             logger.info("Disabling JIT compilation")
             jax.config.update("jax_disable_jit", True)

{xax-0.1.13 → xax-0.1.14}/xax/utils/debugging.py

@@ -53,3 +53,7 @@ def get_named_leaves(
 
 def breakpoint_if_nan(x: Array) -> None:
     jax.lax.cond(jnp.any(jnp.isnan(x)), lambda: jax.debug.breakpoint(), lambda: None)
+
+
+def log_if_nan(x: Array, loc: str) -> None:
+    jax.lax.cond(jnp.any(jnp.isnan(x)), lambda: jax.debug.print("=== NaNs: {loc} ===", loc=loc), lambda: None)

{xax-0.1.13 → xax-0.1.14/xax.egg-info}/PKG-INFO

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