xax 0.1.12__py3-none-any.whl → 0.1.13__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.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/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/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.dist-info → xax-0.1.13.dist-info}/METADATA

@@ -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.dist-info → xax-0.1.13.dist-info}/RECORD

@@ -1,4 +1,4 @@
-xax/__init__.py,sha256=7vdTYO7jAJdDxKZURlFxc3Y5kr5mVQcTQjeh_sYjD6I,13834
+xax/__init__.py,sha256=zBQiNKXLrE-9UPy1q-XBFwPNAKTAxr6wpAwYwaTVggs,13922
 xax/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/requirements-dev.txt,sha256=qkscNkFzWd1S5fump-AKH53rR65v2x5FmboFdy_kKvs,128
 xax/requirements.txt,sha256=9LAEZ5c5gqRSARRVA6xJsVTa4MebPZuC4yOkkwkZJFw,297
@@ -28,7 +28,7 @@ xax/task/loggers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,
 xax/task/loggers/callback.py,sha256=lyuZX6Bir7xJM07ifdQIl1jlclgkiS82UO9V4y7wgPs,1582
 xax/task/loggers/json.py,sha256=yXHb1bmfsEnk4p0F1Up1ertWYdcPAFZm25NT8wE3Jb8,4045
 xax/task/loggers/state.py,sha256=6bG-NRsSUzAukYiglCT0oDj8zRMpffH4e1TKWGw1x4k,959
-xax/task/loggers/stdout.py,sha256=bR0k-PfmFgLfPxLPb4hZw_8G_msA32UeHfAAu11nEYs,6757
+xax/task/loggers/stdout.py,sha256=BBXqr95gNt5KuCN8XyKnTJF8JdwkR4JgLKrkvcaTBVM,6788
 xax/task/loggers/tensorboard.py,sha256=kI8LvBuBBhPgkP8TeaTQb9SQ0FqaIodwQh2SuWDCnIA,7706
 xax/task/mixins/__init__.py,sha256=D3oU31rB9FeOr9MPLleLt5JFbftUr4sBTwgnwQdc2qA,809
 xax/task/mixins/artifacts.py,sha256=2ezmZGzPGe3nhsd9KRkeHWWXdbT9m7drzimIfw6v1XY,2892
@@ -43,7 +43,7 @@ xax/task/mixins/runnable.py,sha256=IYIsLd2k09g-_y6o44EhJqT7E6BpsyEMmsyLSuzqjtc,1
 xax/task/mixins/step_wrapper.py,sha256=-Yu5Nft2CRw1JvZt6J_94SM1vqX8fk08IDK95Pmd2ew,1648
 xax/task/mixins/train.py,sha256=aIebtOIvERYofSyqzNGBpNYlNrXweqFUqM9dHiTx3Dc,26253
 xax/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-xax/utils/debugging.py,sha256=9WlCrEqbq-SVXPEM4rhsLYERH97XNX7XSYLSI3sgKGk,1619
+xax/utils/debugging.py,sha256=0DU41DvYp3SZ9tMrM7sSFpfhC7dieMYR7eRlGNAFrdM,1783
 xax/utils/experiments.py,sha256=5CUja1H_cx4dnVqTGQekOpIhqISwHtAgLxZ34GV7cwM,29229
 xax/utils/jax.py,sha256=tC0NNelbrSTzwNGluiwLGKtoHhVpgdzrv-xherB3VtY,4752
 xax/utils/jaxpr.py,sha256=S80nyEkv188RInzq3kCAdkQCU-bf6s0oPTrCE_LjkRs,2298
@@ -58,8 +58,8 @@ 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=ZCMGfSfr2_b2qZbq9ywPD0zej5tpVSId2JftXpwfB5k,4686
 xax/utils/types/hashable_array.py,sha256=l5iIcFmkYzfGeaZmcSoeFkthFASqM8xJYK3AXhZQYwc,992
-xax-0.1.12.dist-info/licenses/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
-xax-0.1.12.dist-info/METADATA,sha256=hLRAX5__7QjBgjzhxbRftGvEsNrt8IAdgd22dMtHu_Y,1878
-xax-0.1.12.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-xax-0.1.12.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
-xax-0.1.12.dist-info/RECORD,,
+xax-0.1.13.dist-info/licenses/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
+xax-0.1.13.dist-info/METADATA,sha256=P3N5hJMZtXPs199OHPN_cBp57S9zlOGB2B1TRPaczuI,1878
+xax-0.1.13.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
+xax-0.1.13.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
+xax-0.1.13.dist-info/RECORD,,