xax 0.1.12__tar.gz → 0.1.13__tar.gz

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

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

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

@@ -12,7 +12,7 @@ and running the update script:
     python -m scripts.update_api --inplace
 """
 
-__version__ = "0.1.12"
+__version__ = "0.1.13"
 
 # This list shouldn't be modified by hand; instead, run the update script.
 __all__ = [
@@ -75,6 +75,7 @@ __all__ = [
     "Task",
     "collate",
     "collate_non_null",
+    "breakpoint_if_nan",
     "get_named_leaves",
     "BaseFileDownloader",
     "ContextTimer",
@@ -234,6 +235,7 @@ NAME_MAP: dict[str, str] = {
     "Task": "task.task",
     "collate": "utils.data.collate",
     "collate_non_null": "utils.data.collate",
+    "breakpoint_if_nan": "utils.debugging",
     "get_named_leaves": "utils.debugging",
     "BaseFileDownloader": "utils.experiments",
     "ContextTimer": "utils.experiments",
@@ -386,7 +388,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 get_named_leaves
+    from xax.utils.debugging import breakpoint_if_nan, get_named_leaves
     from xax.utils.experiments import (
         BaseFileDownloader,
         ContextTimer,

{xax-0.1.12 → xax-0.1.13}/xax/task/loggers/stdout.py

@@ -14,7 +14,7 @@ from xax.utils.text import Color, colored, format_timedelta
 
 def format_number(value: int | float, precision: int) -> str:
     if isinstance(value, int):
-        return str(value)
+        return f"{value:,}"  # Add commas to the number
     return f"{value:.{precision}g}"
 
 
@@ -80,11 +80,10 @@ class StdoutLogger(LoggerImpl):
         self.write_fp.write("\033[2J\033[H")
 
     def write_state_window(self, line: LogLine) -> None:
-        elapsed_time = format_timedelta(datetime.timedelta(seconds=line.state.elapsed_time_s), short=True)
-        state_info = {
-            "Steps": f"{line.state.num_steps}",
-            "Samples": f"{line.state.num_samples}",
-            "Elapsed Time": f"{elapsed_time}",
+        state_info: dict[str, str] = {
+            "Steps": format_number(line.state.num_steps, 0),
+            "Samples": format_number(line.state.num_samples, 0),
+            "Elapsed Time": format_timedelta(datetime.timedelta(seconds=line.state.elapsed_time_s), short=True),
         }
 
         colored_prefix = colored("Phase: ", "grey", bold=True)

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

@@ -4,6 +4,8 @@ from collections import deque
 from collections.abc import Iterable, Mapping
 from typing import Any, Callable, Deque
 
+import jax
+import jax.numpy as jnp
 from jaxtyping import Array
 
 
@@ -47,3 +49,7 @@ def get_named_leaves(
                     q.append((depth + 1, gname, cnode))
 
     return ret
+
+
+def breakpoint_if_nan(x: Array) -> None:
+    jax.lax.cond(jnp.any(jnp.isnan(x)), lambda: jax.debug.breakpoint(), lambda: None)

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

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