xax 0.1.0__py3-none-any.whl → 0.1.2__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.0"
+__version__ = "0.1.2"
 
 # This list shouldn't be modified by hand; instead, run the update script.
 __all__ = [
@@ -66,6 +66,7 @@ __all__ = [
     "DataloaderConfig",
     "GPUStatsOptions",
     "StepContext",
+    "ValidStepTimer",
     "Script",
     "ScriptConfig",
     "Config",
@@ -219,6 +220,7 @@ NAME_MAP: dict[str, str] = {
     "DataloaderConfig": "task.mixins.data_loader",
     "GPUStatsOptions": "task.mixins.gpu_stats",
     "StepContext": "task.mixins.step_wrapper",
+    "ValidStepTimer": "task.mixins.train",
     "Script": "task.script",
     "ScriptConfig": "task.script",
     "Config": "task.task",
@@ -372,7 +374,7 @@ if IMPORT_ALL or TYPE_CHECKING:
     from xax.task.mixins.data_loader import DataloaderConfig
     from xax.task.mixins.gpu_stats import GPUStatsOptions
     from xax.task.mixins.step_wrapper import StepContext
-    from xax.task.mixins.train import Batch, Output
+    from xax.task.mixins.train import Batch, Output, ValidStepTimer
     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

xax/task/logger.py

@@ -18,7 +18,7 @@ from abc import ABC, abstractmethod
 from collections import defaultdict
 from dataclasses import dataclass
 from types import TracebackType
-from typing import Callable, Iterator, Literal, Self, Sequence, TypeVar, get_args
+from typing import Callable, Iterator, Literal, Self, Sequence, TypeVar, cast, get_args
 
 import jax
 import jax.numpy as jnp
@@ -28,7 +28,7 @@ from PIL import Image, ImageDraw, ImageFont
 from PIL.Image import Image as PILImage
 
 from xax.core.state import Phase, State
-from xax.utils.experiments import IntervalTicker
+from xax.utils.experiments import ContextTimer, IntervalTicker
 from xax.utils.logging import LOG_ERROR_SUMMARY, LOG_PING, LOG_STATUS
 
 logger = logging.getLogger(__name__)
@@ -236,7 +236,7 @@ class LogHistogram:
     num: int
     sum: Number
     sum_squares: Number
-    bucket_limits: list[Number]
+    bucket_limits: list[float]
     bucket_counts: list[int]
 
 
@@ -564,14 +564,14 @@ class Logger:
         Args:
             state: The current step's state.
         """
-        should_log = [logger.should_log(state) for logger in self.loggers]
+        should_log = [lg.should_log(state) for lg in self.loggers]
         if not any(should_log):
             self.clear()
             return
         line = self.pack(state)
         self.clear()
-        for logger in (logger for logger, should_log in zip(self.loggers, should_log) if should_log):
-            logger.write(line)
+        for lg in (lg for lg, should_log in zip(self.loggers, should_log) if should_log):
+            lg.write(line)
 
     def write_error_summary(self, error_summary: str) -> None:
         for logger in self.loggers:
@@ -633,7 +633,10 @@ class Logger:
 
         @functools.lru_cache(maxsize=None)
         def scalar_future() -> Number:
-            return value() if callable(value) else value
+            with ContextTimer() as timer:
+                value_concrete = value() if callable(value) else value
+            logger.debug("Scalar Key: %s, Time: %s", key, timer.elapsed_time)
+            return value_concrete
 
         self.scalars[namespace][key] = scalar_future
 
@@ -657,7 +660,9 @@ class Logger:
 
         @functools.lru_cache(maxsize=None)
         def distribution_future() -> LogDistribution:
-            mean, std = value() if callable(value) else value
+            with ContextTimer() as timer:
+                mean, std = value() if callable(value) else value
+            logger.debug("Distribution Key: %s, Time: %s", key, timer.elapsed_time)
             return LogDistribution(mean=mean, std=std)
 
         self.distributions[namespace][key] = distribution_future
@@ -684,26 +689,84 @@ class Logger:
 
         @functools.lru_cache(maxsize=None)
         def histogram_future() -> LogHistogram:
-            values = value() if callable(value) else value
-            values = values.reshape(-1)  # Must be flat.
-
-            if isinstance(values, Array):
-                counts, limits = jnp.histogram(values, bins=bins)
-                counts, limits = as_numpy(counts), as_numpy(limits)
-            elif isinstance(values, np.ndarray):
-                counts, limits = np.histogram(values, bins=bins)
-            else:
-                raise ValueError(f"Unsupported histogram type: {type(values)}")
-
-            return LogHistogram(
-                min=float(values.min()),
-                max=float(values.max()),
-                num=int(values.size),
-                sum=float(values.sum()),
-                sum_squares=float(values.dot(values)),
-                bucket_limits=limits[1:].tolist(),
-                bucket_counts=counts.tolist(),
-            )
+            with ContextTimer() as timer:
+                values = value() if callable(value) else value
+                values = values.reshape(-1)  # Must be flat.
+
+                if isinstance(values, Array):
+                    counts, limits = jnp.histogram(values, bins=bins)
+                    counts, limits = as_numpy(counts), as_numpy(limits)
+                elif isinstance(values, np.ndarray):
+                    counts, limits = np.histogram(values, bins=bins)
+                else:
+                    raise ValueError(f"Unsupported histogram type: {type(values)}")
+
+                histogram_values = LogHistogram(
+                    min=float(values.min()),
+                    max=float(values.max()),
+                    num=int(values.size),
+                    sum=float(values.sum()),
+                    sum_squares=float(values.dot(values)),
+                    bucket_limits=cast(list[float], limits[1:].tolist()),
+                    bucket_counts=cast(list[int], counts.tolist()),
+                )
+
+            logger.debug("Histogram Key: %s, Time: %s", key, timer.elapsed_time)
+            return histogram_values
+
+        self.histograms[namespace][key] = histogram_future
+
+    def log_histogram_raw(
+        self,
+        key: str,
+        counts: Array | np.ndarray,
+        limits: Array | np.ndarray,
+        minv: Number | None = None,
+        maxv: Number | None = None,
+        sumv: Number | None = None,
+        sum_squaresv: Number | None = None,
+        *,
+        namespace: str | None = None,
+    ) -> None:
+        """Logs a histogram from raw counts and limits.
+
+        Args:
+            key: The key being logged
+            counts: The counts of the histogram
+            limits: The limits of the histogram
+            minv: The minimum value of the histogram
+            maxv: The maximum value of the histogram
+            sumv: The sum of the histogram
+            sum_squaresv: The sum of the squares of the histogram
+            namespace: An optional logging namespace
+        """
+        if not self.active:
+            raise RuntimeError("The logger is not active")
+        namespace = self.resolve_namespace(namespace)
+
+        @functools.lru_cache(maxsize=None)
+        def histogram_future() -> LogHistogram:
+            with ContextTimer() as timer:
+                counts_np = (as_numpy(counts) if isinstance(counts, Array) else counts).astype(int)
+                limits_np = (as_numpy(limits) if isinstance(limits, Array) else limits).astype(float)
+
+                minv_ = counts_np.min() if minv is None else minv
+                maxv_ = counts_np.max() if maxv is None else maxv
+                sumv_ = counts_np.sum() if sumv is None else sumv
+                sum_squaresv_ = counts_np.dot(counts_np) if sum_squaresv is None else sum_squaresv
+
+                histogram_values = LogHistogram(
+                    min=float(minv_),
+                    max=float(maxv_),
+                    num=int(counts_np.size),
+                    sum=float(sumv_),
+                    sum_squares=float(sum_squaresv_),
+                    bucket_limits=cast(list[float], limits_np.tolist()),
+                    bucket_counts=cast(list[int], counts_np.tolist()),
+                )
+
+            logger.debug("Raw Histogram Key: %s, Time: %s", key, timer.elapsed_time)
+            return histogram_values
 
         self.histograms[namespace][key] = histogram_future
 
@@ -748,7 +811,10 @@ class Logger:
 
         @functools.lru_cache(maxsize=None)
         def image_future() -> LogImage:
-            return get_image(value() if callable(value) else value, target_resolution)
+            with ContextTimer() as timer:
+                image = get_image(value() if callable(value) else value, target_resolution)
+            logger.debug("Image Key: %s, Time: %s", key, timer.elapsed_time)
+            return image
 
         self.images[namespace][key] = image_future
 
@@ -784,15 +850,20 @@ class Logger:
 
         @functools.lru_cache(maxsize=None)
         def image_future() -> LogImage:
-            image, label = value() if callable(value) else value
-            image = get_image(image, target_resolution)
-            return image_with_text(
-                image.image,
-                standardize_text(label, max_line_length),
-                max_num_lines=max_num_lines,
-                line_spacing=line_spacing,
-                centered=centered,
-            )
+            with ContextTimer() as timer:
+                image, label = value() if callable(value) else value
+                image = get_image(image, target_resolution)
+
+                image_value = image_with_text(
+                    image.image,
+                    standardize_text(label, max_line_length),
+                    max_num_lines=max_num_lines,
+                    line_spacing=line_spacing,
+                    centered=centered,
+                )
+
+            logger.debug("Labeled Image Key: %s, Time: %s", key, timer.elapsed_time)
+            return image_value
 
         self.images[namespace][key] = image_future
 
@@ -831,15 +902,18 @@ class Logger:
 
         @functools.lru_cache(maxsize=None)
         def images_future() -> LogImage:
-            images = value() if callable(value) else value
-            if max_images is not None:
-                images = images[:max_images]
-            if isinstance(images, Array):
-                images = as_numpy(images)
-            if isinstance(images, Sequence):
-                images = list(images)
-            images = [get_image(image, target_resolution) for image in images]
-            tiled = tile_images([img.image for img in images], sep)
+            with ContextTimer() as timer:
+                images = value() if callable(value) else value
+                if max_images is not None:
+                    images = images[:max_images]
+                if isinstance(images, Array):
+                    images = as_numpy(images)
+                if isinstance(images, Sequence):
+                    images = list(images)
+                images = [get_image(image, target_resolution) for image in images]
+                tiled = tile_images([img.image for img in images], sep)
+
+            logger.debug("Images Key: %s, Time: %s", key, timer.elapsed_time)
             return LogImage(image=tiled)
 
         self.images[namespace][key] = images_future
@@ -886,22 +960,25 @@ class Logger:
 
         @functools.lru_cache(maxsize=None)
         def images_future() -> LogImage:
-            images, labels = value() if callable(value) else value
-            if max_images is not None:
-                images = images[:max_images]
-                labels = labels[:max_images]
-            images = [get_image(image, target_resolution) for image in images]
-            labeled = [
-                image_with_text(
-                    img.image,
-                    standardize_text(label, max_line_length),
-                    max_num_lines=max_num_lines,
-                    line_spacing=line_spacing,
-                    centered=centered,
-                )
-                for img, label in zip(images, labels)
-            ]
-            tiled = tile_images([img.image for img in labeled], sep)
+            with ContextTimer() as timer:
+                images, labels = value() if callable(value) else value
+                if max_images is not None:
+                    images = images[:max_images]
+                    labels = labels[:max_images]
+                images = [get_image(image, target_resolution) for image in images]
+                labeled = [
+                    image_with_text(
+                        img.image,
+                        standardize_text(label, max_line_length),
+                        max_num_lines=max_num_lines,
+                        line_spacing=line_spacing,
+                        centered=centered,
+                    )
+                    for img, label in zip(images, labels)
+                ]
+                tiled = tile_images([img.image for img in labeled], sep)
+
+            logger.debug("Labeled Images Key: %s, Time: %s", key, timer.elapsed_time)
             return LogImage(image=tiled)
 
         self.images[namespace][key] = images_future
@@ -934,7 +1011,11 @@ class Logger:
 
         @functools.lru_cache(maxsize=None)
         def video_future() -> LogVideo:
-            return get_video(value() if callable(value) else value, fps=fps)
+            with ContextTimer() as timer:
+                video = get_video(value() if callable(value) else value, fps=fps)
+
+            logger.debug("Video Key: %s, Time: %s", key, timer.elapsed_time)
+            return video
 
         self.videos[namespace][key] = video_future
 

xax/task/mixins/train.py

@@ -156,7 +156,6 @@ class TrainConfig(
     valid_first_n_steps: int = field(0, help="Treat the first N steps as validation steps")
     valid_every_n_seconds: float | None = field(60.0 * 10.0, help="Run validation every N seconds")
     valid_first_n_seconds: float | None = field(60.0, help="Run first validation after N seconds")
-    batch_dim: int = field(0, help="The batch dimension, for splitting batches into chunks")
     max_steps: int | None = field(None, help="Maximum number of steps to run")
     step_kind: str = field("step", help=f"How to measure a step; one of [{', '.join(get_args(StepKind))}]")
     random_seed: int = field(1337, help="Random seed for the task")

xax/utils/experiments.py

@@ -109,8 +109,8 @@ class StateTimer:
 
     def step(self, state: State) -> None:
         cur_time = time.time()
-        self.step_timer.step(state.num_steps, cur_time)
-        self.sample_timer.step(state.num_samples, cur_time)
+        self.step_timer.step(state.num_steps if state.phase == "train" else state.num_valid_steps, cur_time)
+        self.sample_timer.step(state.num_samples if state.phase == "train" else state.num_valid_samples, cur_time)
         self.iter_timer.step(cur_time)
 
     def log_dict(self) -> dict[str, dict[str, int | float]]:

xax/utils/logging.py

@@ -140,7 +140,13 @@ class ColoredFormatter(logging.Formatter):
         return logging.Formatter.format(self, record)
 
 
-def configure_logging(prefix: str | None = None, *, rank: int | None = None, world_size: int | None = None) -> None:
+def configure_logging(
+    prefix: str | None = None,
+    *,
+    rank: int | None = None,
+    world_size: int | None = None,
+    debug: bool | None = None,
+) -> None:
     """Instantiates logging.
 
     This captures logs and reroutes them to the Toasts module, which is
@@ -151,6 +157,7 @@ def configure_logging(prefix: str | None = None, *, rank: int | None = None, wor
         prefix: An optional prefix to add to the logger
         rank: The current rank, or None if not using multiprocessing
         world_size: The total world size, or None if not using multiprocessing
+        debug: Whether to enable debug logging
     """
     if rank is not None or world_size is not None:
         assert rank is not None and world_size is not None
@@ -168,7 +175,10 @@ def configure_logging(prefix: str | None = None, *, rank: int | None = None, wor
     stream_handler.addFilter(filter)
     root_logger.addHandler(stream_handler)
 
-    root_logger.setLevel(logging._nameToLevel[config.log_level])
+    if debug is None:
+        root_logger.setLevel(logging._nameToLevel[config.log_level])
+    else:
+        root_logger.setLevel(logging.DEBUG if debug else logging.INFO)
 
     # Avoid junk logs from other libraries.
     if config.hide_third_party_logs:

{xax-0.1.0.dist-info → xax-0.1.2.dist-info}/METADATA

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

{xax-0.1.0.dist-info → xax-0.1.2.dist-info}/RECORD

@@ -1,4 +1,4 @@
-xax/__init__.py,sha256=psnt49vcnlzZ7llojgBqCEP0ZquHws_8tZpAZ-5vvLE,13280
+xax/__init__.py,sha256=Ti6hrfoY5wnywzOvkvtCwq2SvLsjYfbm_6U_UzYakls,13361
 xax/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/requirements-dev.txt,sha256=qkscNkFzWd1S5fump-AKH53rR65v2x5FmboFdy_kKvs,128
 xax/requirements.txt,sha256=9LAEZ5c5gqRSARRVA6xJsVTa4MebPZuC4yOkkwkZJFw,297
@@ -15,7 +15,7 @@ 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/logger.py,sha256=z5tcFWsoQkDExCZSL7k_ub-a2RfxmGGGn0WOnYBG82Y,32631
+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
 xax/task/launchers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -39,13 +39,13 @@ xax/task/mixins/logger.py,sha256=6oXsJJyNUx6YT3q58FVXMZBUpMgjVkGre6BXFN20cVI,280
 xax/task/mixins/process.py,sha256=d1opVgvc6bOFXb7R58b07F4P5lbSZIzYaajtE0eBbpw,1477
 xax/task/mixins/runnable.py,sha256=IYIsLd2k09g-_y6o44EhJqT7E6BpsyEMmsyLSuzqjtc,1979
 xax/task/mixins/step_wrapper.py,sha256=-Yu5Nft2CRw1JvZt6J_94SM1vqX8fk08IDK95Pmd2ew,1648
-xax/task/mixins/train.py,sha256=6IW1gNnE1a92E1h6SThmnGS9dyjqw0KrWzx0SG6n0_Q,22318
+xax/task/mixins/train.py,sha256=BEC7HSwBlGZDe7jCsedqEA8-K1Zx52-bTjsBONYIE5g,22225
 xax/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/utils/debugging.py,sha256=9WlCrEqbq-SVXPEM4rhsLYERH97XNX7XSYLSI3sgKGk,1619
-xax/utils/experiments.py,sha256=_cwoBaiBxoQ_Tstm0rz7TEqfELqcktmPflb6AP1K0qA,28779
+xax/utils/experiments.py,sha256=d-e-RCw9PlnuqV3FPW0U74zcvlOKV48lUrX8tvAfhew,28887
 xax/utils/jax.py,sha256=tC0NNelbrSTzwNGluiwLGKtoHhVpgdzrv-xherB3VtY,4752
 xax/utils/jaxpr.py,sha256=S80nyEkv188RInzq3kCAdkQCU-bf6s0oPTrCE_LjkRs,2298
-xax/utils/logging.py,sha256=ST1hp2C2xntVVJBUHwo3YxPK19fBLNvHU2WGO1xqcXA,6418
+xax/utils/logging.py,sha256=GAhTne2rdB4Fa1lzk06DMO15U8MTejn6XTClShC-ZtU,6622
 xax/utils/numpy.py,sha256=_jOXVi-d2AtJnRftPkRK5MDMzsU8slgw-Jjv4GRm6ns,1197
 xax/utils/profile.py,sha256=-aFdWpgYFvBsBZXSLL4zXrFe3zzsDqzmx4q5f2WOtpQ,1628
 xax/utils/pytree.py,sha256=7GjQoPc_ZSZt3QS_9qXoBWl1jfMp1qZa7aViQoWJ0OQ,8864
@@ -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.0.dist-info/licenses/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
-xax-0.1.0.dist-info/METADATA,sha256=b5q3AVoywNcDoTLcsf5mCk3wSmL4BHDAudewxmY1XJw,1877
-xax-0.1.0.dist-info/WHEEL,sha256=1tXe9gY0PYatrMPMDd6jXqjfpz_B-Wqm32CPfRC58XU,91
-xax-0.1.0.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
-xax-0.1.0.dist-info/RECORD,,
+xax-0.1.2.dist-info/licenses/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
+xax-0.1.2.dist-info/METADATA,sha256=-BB6_Qiip_pPkf96Wl9FZsM_7PPKJr5l8v2owrXCvoI,1877
+xax-0.1.2.dist-info/WHEEL,sha256=L0N565qmK-3nM2eBoMNFszYJ_MTx03_tQ0CQu1bHLYo,91
+xax-0.1.2.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
+xax-0.1.2.dist-info/RECORD,,

{xax-0.1.0.dist-info → xax-0.1.2.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (77.0.3)
+Generator: setuptools (78.0.1)
 Root-Is-Purelib: true
 Tag: py3-none-any