xax 0.0.5__py3-none-any.whl → 0.0.6__py3-none-any.whl

xax/task/logger.py

@@ -24,7 +24,6 @@ import jax
 import jax.numpy as jnp
 import numpy as np
 from jaxtyping import Array
-from omegaconf import DictConfig
 from PIL import Image, ImageDraw, ImageFont
 from PIL.Image import Image as PILImage
 
@@ -92,50 +91,6 @@ def make_human_viewable_resolution(
     return image.resize((new_width, new_height), interpolation)
 
 
-def image_with_text(
-    image: PILImage,
-    text: list[str],
-    max_num_lines: int | None,
-    line_spacing: int,
-    centered: bool,
-) -> PILImage:
-    """Adds a text label to an image.
-
-    Args:
-        image: The image to label, with shape (C, H, W)
-        text: The text label for the image
-        max_num_lines: The number of lines of spacing to add to the bottom
-            of the image
-        line_spacing: The spacing between adjacent lines
-        centered: If set, center the text labels, otherwise align to the left
-
-    Returns:
-        The image with a text label
-    """
-    if not text:
-        return image
-    if max_num_lines is None:
-        max_num_lines = len(text)
-    else:
-        text = text[:max_num_lines]
-    width, height = image.size
-    font: ImageFont.ImageFont = ImageFont.load_default()
-    _, _, _, line_height = font.getbbox(text[0])
-    new_width, new_height = width, height + line_spacing + max_num_lines * (line_height + line_spacing)
-    padded_image = Image.new(image.mode, (new_width, new_height), 255)
-    padded_image.paste(image, (0, 0))
-    drawer = ImageDraw.Draw(padded_image)
-    for i, text_line in enumerate(text):
-        text_line_top = height + line_spacing + i * (line_height + line_spacing)
-        if centered:
-            _, _, line_width, _ = font.getbbox(text_line)
-            text_line_left = (width - line_width) / 2
-            drawer.text((text_line_left, text_line_top), text_line, font=font, fill=0)
-        else:
-            drawer.text((line_spacing, text_line_top), text_line, font=font, fill=0)
-    return padded_image
-
-
 class namespace_context:  # noqa: N801
     def __init__(self, name: str | None) -> None:
         self._name = name
@@ -250,7 +205,68 @@ def as_numpy(array: Array) -> np.ndarray:
     return np.array(array)
 
 
-def get_image(image: np.ndarray | Array | PILImage, target_resolution: tuple[int, int] | None = None) -> PILImage:
+@dataclass(kw_only=True)
+class LogImage:
+    image: PILImage
+
+
+@dataclass(kw_only=True)
+class LogVideo:
+    """Container for video data and metadata.
+
+    Attributes:
+        frames: Video frames as a numpy array of shape (T,H,W,C)
+        fps: Frames per second
+    """
+
+    frames: np.ndarray
+    fps: int
+
+
+@dataclass(kw_only=True)
+class LogLine:
+    state: State
+    scalars: dict[str, dict[str, Number]]
+    strings: dict[str, dict[str, str]]
+    images: dict[str, dict[str, LogImage]]
+    videos: dict[str, dict[str, LogVideo]]
+
+
+@dataclass(kw_only=True)
+class LogErrorSummary:
+    message: str
+
+
+@dataclass(kw_only=True)
+class LogError:
+    message: str
+    location: str | None = None
+
+    @property
+    def message_with_location(self) -> str:
+        message = self.message
+        if self.location is not None:
+            message += f" ({self.location})"
+        return message
+
+
+@dataclass(kw_only=True)
+class LogStatus:
+    message: str
+    created: float
+    filename: str | None = None
+    lineno: int | None = None
+
+
+@dataclass(kw_only=True)
+class LogPing:
+    message: str
+    created: float
+    filename: str | None = None
+    lineno: int | None = None
+
+
+def get_image(image: np.ndarray | Array | PILImage, target_resolution: tuple[int, int] | None = None) -> LogImage:
     if not isinstance(image, (np.ndarray, Array, PILImage)):
         raise ValueError(f"Unsupported image type: {type(image)}")
     if isinstance(image, Array):
@@ -283,54 +299,85 @@ def get_image(image: np.ndarray | Array | PILImage, target_resolution: tuple[int
 
     if target_resolution is not None:
         image = make_human_viewable_resolution(image, trg_res=target_resolution)
-    return image
+    return LogImage(image=image)
 
 
-@dataclass
-class LogImage:
-    image: PILImage
+def image_with_text(
+    image: PILImage,
+    text: list[str],
+    max_num_lines: int | None,
+    line_spacing: int,
+    centered: bool,
+) -> LogImage:
+    """Adds a text label to an image.
 
+    Args:
+        image: The image to label, with shape (C, H, W)
+        text: The text label for the image
+        max_num_lines: The number of lines of spacing to add to the bottom
+            of the image
+        line_spacing: The spacing between adjacent lines
+        centered: If set, center the text labels, otherwise align to the left
 
-@dataclass
-class LogLine:
-    state: State
-    scalars: dict[str, dict[str, Number]]
-    strings: dict[str, dict[str, str]]
-    images: dict[str, dict[str, LogImage]]
+    Returns:
+        The image with a text label
+    """
+    if not text:
+        return LogImage(image=image)
+    if max_num_lines is None:
+        max_num_lines = len(text)
+    else:
+        text = text[:max_num_lines]
+    width, height = image.size
+    font: ImageFont.ImageFont = ImageFont.load_default()
+    _, _, _, line_height = font.getbbox(text[0])
+    new_width, new_height = width, height + line_spacing + max_num_lines * (line_height + line_spacing)
+    padded_image = Image.new(image.mode, (new_width, new_height), 255)
+    padded_image.paste(image, (0, 0))
+    drawer = ImageDraw.Draw(padded_image)
+    for i, text_line in enumerate(text):
+        text_line_top = height + line_spacing + i * (line_height + line_spacing)
+        if centered:
+            _, _, line_width, _ = font.getbbox(text_line)
+            text_line_left = (width - line_width) / 2
+            drawer.text((text_line_left, text_line_top), text_line, font=font, fill=0)
+        else:
+            drawer.text((line_spacing, text_line_top), text_line, font=font, fill=0)
+    return LogImage(image=padded_image)
 
 
-@dataclass
-class LogErrorSummary:
-    message: str
+def get_video(video: np.ndarray | Array, fps: int = 30) -> LogVideo:
+    """Converts video data to standard format.
 
+    Args:
+        video: The video frames. Can be:
+            - A numpy array of shape (T, H, W, C) or (T, C, H, W)
+            - A JAX array of shape (T, H, W, C) or (T, C, H, W)
+        fps: Frames per second
 
-@dataclass
-class LogError:
-    message: str
-    location: str | None = None
+    Returns:
+        LogVideo containing standardized video frames
+    """
+    if isinstance(video, Array):
+        video = as_numpy(video)
 
-    @property
-    def message_with_location(self) -> str:
-        message = self.message
-        if self.location is not None:
-            message += f" ({self.location})"
-        return message
+    if not isinstance(video, np.ndarray):
+        raise ValueError(f"Unsupported video type: {type(video)}")
 
+    # Handle different dimension orderings
+    if video.ndim != 4:
+        raise ValueError(f"Expected video array of shape (T, H, W, C) or (T, C, H, W), got shape {video.shape}")
 
-@dataclass
-class LogStatus:
-    message: str
-    created: float
-    filename: str | None = None
-    lineno: int | None = None
+    if video.shape[1] == 3:  # (T,C,H,W) format
+        video = video.transpose(0, 2, 3, 1)
 
+    # Normalize and convert to uint8 if needed
+    if np.issubdtype(video.dtype, np.floating):
+        video = (normalize(video) * 255).round().astype(np.uint8)
+    elif video.dtype != np.uint8:
+        raise ValueError(f"Unsupported video dtype: {video.dtype}")
 
-@dataclass
-class LogPing:
-    message: str
-    created: float
-    filename: str | None = None
-    lineno: int | None = None
+    return LogVideo(frames=video, fps=fps)
 
 
 class LoggerImpl(ABC):
@@ -391,25 +438,12 @@ class LoggerImpl(ABC):
             ping: The ping to write.
         """
 
-    def log_git_state(self, git_state: str) -> None:
-        """Logs Git state for the current run.
-
-        Args:
-            git_state: The Git state, as text blocks.
-        """
-
-    def log_training_code(self, training_code: str) -> None:
-        """Logs the training script code.
-
-        Args:
-            training_code: The training script code.
-        """
-
-    def log_config(self, config: DictConfig) -> None:
-        """Logs the configuration for the current run.
+    def log_file(self, name: str, contents: str) -> None:
+        """Logs a large text file.
 
         Args:
-            config: The configuration, as a DictConfig.
+            name: The name of the file.
+            contents: The contents of the file.
         """
 
     def should_log(self, state: State) -> bool:
@@ -464,7 +498,8 @@ class Logger:
     def __init__(self, default_namespace: str = DEFAULT_NAMESPACE) -> None:
         self.scalars: dict[str, dict[str, Callable[[], Number]]] = defaultdict(dict)
         self.strings: dict[str, dict[str, Callable[[], str]]] = defaultdict(dict)
-        self.images: dict[str, dict[str, Callable[[], PILImage]]] = defaultdict(dict)
+        self.images: dict[str, dict[str, Callable[[], LogImage]]] = defaultdict(dict)
+        self.videos: dict[str, dict[str, Callable[[], LogVideo]]] = defaultdict(dict)
         self.default_namespace = default_namespace
         self.loggers: list[LoggerImpl] = []
 
@@ -488,13 +523,15 @@ class Logger:
             state=state,
             scalars={k: {kk: v() for kk, v in v.items()} for k, v in self.scalars.items()},
             strings={k: {kk: v() for kk, v in v.items()} for k, v in self.strings.items()},
-            images={k: {kk: LogImage(v()) for kk, v in v.items()} for k, v in self.images.items()},
+            images={k: {kk: v() for kk, v in v.items()} for k, v in self.images.items()},
+            videos={k: {kk: v() for kk, v in v.items()} for k, v in self.videos.items()},
         )
 
     def clear(self) -> None:
         self.scalars.clear()
         self.strings.clear()
         self.images.clear()
+        self.videos.clear()
 
     def write(self, state: State) -> None:
         """Writes the current step's logging information.
@@ -513,11 +550,11 @@ class Logger:
 
     def write_error_summary(self, error_summary: str) -> None:
         for logger in self.loggers:
-            logger.write_error_summary(LogErrorSummary(error_summary))
+            logger.write_error_summary(LogErrorSummary(message=error_summary))
 
     def write_error(self, message: str, location: str | None = None) -> None:
         for logger in self.loggers:
-            logger.write_error(LogError(message, location))
+            logger.write_error(LogError(message=message, location=location))
 
     def write_status(
         self,
@@ -526,7 +563,12 @@ class Logger:
         lineno: int | None = None,
         created: float | None = None,
     ) -> None:
-        status = LogStatus(message, time.time() if created is None else created, filename, lineno)
+        status = LogStatus(
+            message=message,
+            created=time.time() if created is None else created,
+            filename=filename,
+            lineno=lineno,
+        )
         for logger in self.loggers:
             logger.write_status(status)
 
@@ -537,7 +579,12 @@ class Logger:
         lineno: int | None = None,
         created: float | None = None,
     ) -> None:
-        ping = LogPing(message, time.time() if created is None else created, filename, lineno)
+        ping = LogPing(
+            message=message,
+            created=time.time() if created is None else created,
+            filename=filename,
+            lineno=lineno,
+        )
         for logger in self.loggers:
             logger.write_ping(ping)
 
@@ -556,6 +603,9 @@ class Logger:
             raise RuntimeError("The logger is not active")
         namespace = self.resolve_namespace(namespace)
 
+        if isinstance(value, jnp.ndarray):
+            assert value.ndim == 0, f"Scalar must be a 0D array, got shape {value.shape}"
+
         @functools.lru_cache(maxsize=None)
         def scalar_future() -> Number:
             return value() if callable(value) else value
@@ -602,7 +652,7 @@ class Logger:
         namespace = self.resolve_namespace(namespace)
 
         @functools.lru_cache(maxsize=None)
-        def image_future() -> PILImage:
+        def image_future() -> LogImage:
             return get_image(value() if callable(value) else value, target_resolution)
 
         self.images[namespace][key] = image_future
@@ -638,11 +688,11 @@ class Logger:
         namespace = self.resolve_namespace(namespace)
 
         @functools.lru_cache(maxsize=None)
-        def image_future() -> PILImage:
+        def image_future() -> LogImage:
             image, label = value() if callable(value) else value
             image = get_image(image, target_resolution)
             return image_with_text(
-                image,
+                image.image,
                 standardize_text(label, max_line_length),
                 max_num_lines=max_num_lines,
                 line_spacing=line_spacing,
@@ -685,7 +735,7 @@ class Logger:
         namespace = self.resolve_namespace(namespace)
 
         @functools.lru_cache(maxsize=None)
-        def images_future() -> PILImage:
+        def images_future() -> LogImage:
             images = value() if callable(value) else value
             if max_images is not None:
                 images = images[:max_images]
@@ -694,7 +744,8 @@ class Logger:
             if isinstance(images, Sequence):
                 images = list(images)
             images = [get_image(image, target_resolution) for image in images]
-            return tile_images(images, sep)
+            tiled = tile_images([img.image for img in images], sep)
+            return LogImage(image=tiled)
 
         self.images[namespace][key] = images_future
 
@@ -739,37 +790,58 @@ class Logger:
         namespace = self.resolve_namespace(namespace)
 
         @functools.lru_cache(maxsize=None)
-        def images_future() -> PILImage:
+        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]
-            images = [
+            labeled = [
                 image_with_text(
-                    image,
+                    img.image,
                     standardize_text(label, max_line_length),
                     max_num_lines=max_num_lines,
                     line_spacing=line_spacing,
                     centered=centered,
                 )
-                for image, label in zip(images, labels)
+                for img, label in zip(images, labels)
             ]
-            return tile_images(images, sep)
+            tiled = tile_images([img.image for img in labeled], sep)
+            return LogImage(image=tiled)
 
         self.images[namespace][key] = images_future
 
-    def log_git_state(self, git_state: str) -> None:
+    def log_file(self, name: str, contents: str) -> None:
         for logger in self.loggers:
-            logger.log_git_state(git_state)
+            logger.log_file(name, contents)
 
-    def log_training_code(self, training_code: str) -> None:
-        for logger in self.loggers:
-            logger.log_training_code(training_code)
+    def log_video(
+        self,
+        key: str,
+        value: Callable[[], np.ndarray | Array] | np.ndarray | Array,
+        *,
+        fps: int = 30,
+        namespace: str | None = None,
+    ) -> None:
+        """Logs a video.
 
-    def log_config(self, config: DictConfig) -> None:
-        for logger in self.loggers:
-            logger.log_config(config)
+        Args:
+            key: The key being logged
+            value: The video frames. Can be:
+                - A numpy array of shape (T,H,W,C) or (T,C,H,W)
+                - A JAX array of shape (T,H,W,C) or (T,C,H,W)
+            fps: Frames per second
+            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 video_future() -> LogVideo:
+            return get_video(value() if callable(value) else value, fps=fps)
+
+        self.videos[namespace][key] = video_future
 
     def __enter__(self) -> Self:
         self.active = True

xax/task/loggers/callback.py

@@ -2,8 +2,6 @@
 
 from typing import Callable
 
-from omegaconf import DictConfig
-
 from xax.task.logger import LogError, LogErrorSummary, LoggerImpl, LogLine, LogPing, LogStatus
 
 
@@ -16,9 +14,7 @@ class CallbackLogger(LoggerImpl):
         error_callback: Callable[[LogError], None] = lambda x: None,
         status_callback: Callable[[LogStatus], None] = lambda x: None,
         ping_callback: Callable[[LogPing], None] = lambda x: None,
-        git_state_callback: Callable[[str], None] = lambda x: None,
-        training_code_callback: Callable[[str], None] = lambda x: None,
-        config_callback: Callable[[DictConfig], None] = lambda x: None,
+        file_callback: Callable[[str, str], None] = lambda x, y: None,
     ) -> None:
         super().__init__()
 
@@ -27,9 +23,7 @@ class CallbackLogger(LoggerImpl):
         self.error_callback = error_callback
         self.status_callback = status_callback
         self.ping_callback = ping_callback
-        self.git_state_callback = git_state_callback
-        self.training_code_callback = training_code_callback
-        self.config_callback = config_callback
+        self.file_callback = file_callback
 
     def write(self, line: LogLine) -> None:
         self.callback(line)
@@ -46,11 +40,5 @@ class CallbackLogger(LoggerImpl):
     def write_ping(self, ping: LogPing) -> None:
         self.ping_callback(ping)
 
-    def log_git_state(self, git_state: str) -> None:
-        self.git_state_callback(git_state)
-
-    def log_training_code(self, training_code: str) -> None:
-        self.training_code_callback(training_code)
-
-    def log_config(self, config: DictConfig) -> None:
-        self.config_callback(config)
+    def log_file(self, name: str, contents: str) -> None:
+        self.file_callback(name, contents)

xax/task/loggers/state.py

@@ -3,8 +3,6 @@
 from pathlib import Path
 from typing import Literal
 
-from omegaconf import DictConfig, OmegaConf
-
 from xax.task.logger import LoggerImpl, LogLine
 
 
@@ -12,9 +10,6 @@ class StateLogger(LoggerImpl):
     def __init__(
         self,
         run_directory: str | Path,
-        git_state_name: str = "git_state.txt",
-        train_code_name: str = "train_code.py",
-        config_name: str = "config.yaml",
         flush_immediately: bool = False,
         open_mode: Literal["w", "a"] = "w",
         line_sep: str = "\n",
@@ -22,24 +17,16 @@ class StateLogger(LoggerImpl):
     ) -> None:
         super().__init__(float("inf"))
 
-        self.git_state_file = Path(run_directory).expanduser().resolve() / git_state_name
-        self.train_code_file = Path(run_directory).expanduser().resolve() / train_code_name
-        self.config_file = Path(run_directory).expanduser().resolve() / config_name
+        self.run_directory = Path(run_directory).expanduser().resolve()
+
         self.flush_immediately = flush_immediately
         self.open_mode = open_mode
         self.line_sep = line_sep
         self.remove_unicode_from_namespaces = remove_unicode_from_namespaces
 
-    def log_git_state(self, git_state: str) -> None:
-        with open(self.git_state_file, "w") as f:
-            f.write(git_state)
-
-    def log_training_code(self, training_code: str) -> None:
-        with open(self.train_code_file, "w") as f:
-            f.write(training_code)
-
-    def log_config(self, config: DictConfig) -> None:
-        OmegaConf.save(config, self.config_file)
+    def log_file(self, name: str, contents: str) -> None:
+        with open(self.run_directory / name, "w") as f:
+            f.write(contents)
 
     def write(self, line: LogLine) -> None:
         pass

xax/task/loggers/tensorboard.py

@@ -12,8 +12,6 @@ import time
 from pathlib import Path
 from typing import TypeVar
 
-from omegaconf import DictConfig, OmegaConf
-
 from xax.core.state import Phase
 from xax.nn.parallel import is_master
 from xax.task.logger import LoggerImpl, LogLine
@@ -60,10 +58,7 @@ class TensorboardLogger(LoggerImpl):
 
         self.proc: subprocess.Popen | None = None
 
-        self.git_state: str | None = None
-        self.training_code: str | None = None
-        self.config: DictConfig | None = None
-
+        self.files: dict[str, str] = {}
         self.writers = TensorboardWriters(log_directory=self.log_directory, flush_seconds=flush_seconds)
         self._started = False
 
@@ -158,20 +153,10 @@ class TensorboardLogger(LoggerImpl):
         self._start()
         return self.writers.writer(phase)
 
-    def log_git_state(self, git_state: str) -> None:
-        if not is_master():
-            return
-        self.git_state = f"```\n{git_state}\n```"
-
-    def log_training_code(self, training_code: str) -> None:
+    def log_file(self, name: str, contents: str) -> None:
         if not is_master():
             return
-        self.training_code = f"```python\n{training_code}\n```"
-
-    def log_config(self, config: DictConfig) -> None:
-        if not is_master():
-            return
-        self.config = config
+        self.files[name] = f"```\n{contents}\n```"
 
     def write(self, line: LogLine) -> None:
         if not is_master():
@@ -210,14 +195,15 @@ class TensorboardLogger(LoggerImpl):
                     walltime=walltime,
                 )
 
-        if self.config is not None:
-            writer.add_text("config", f"```\n{OmegaConf.to_yaml(self.config)}\n```")
-            self.config = None
-
-        if self.git_state is not None:
-            writer.add_text("git", self.git_state)
-            self.git_state = None
+        for namespace, videos in line.videos.items():
+            for video_key, video_value in videos.items():
+                writer.add_video(
+                    f"{namespace}/{video_key}",
+                    video_value.frames,
+                    fps=video_value.fps,
+                    global_step=line.state.num_steps,
+                )
 
-        if self.training_code is not None:
-            writer.add_text("code", self.training_code)
-            self.training_code = None
+        for name, contents in self.files.items():
+            writer.add_text(name, contents)
+        self.files.clear()

xax/task/mixins/__init__.py

@@ -2,6 +2,7 @@
 
 from xax.task.mixins.artifacts import ArtifactsConfig, ArtifactsMixin
 from xax.task.mixins.checkpointing import CheckpointingConfig, CheckpointingMixin
+from xax.task.mixins.compile import CompileConfig, CompileMixin
 from xax.task.mixins.cpu_stats import CPUStatsConfig, CPUStatsMixin
 from xax.task.mixins.data_loader import DataloadersConfig, DataloadersMixin
 from xax.task.mixins.gpu_stats import GPUStatsConfig, GPUStatsMixin

xax/task/mixins/artifacts.py

@@ -8,6 +8,8 @@ from dataclasses import dataclass
 from pathlib import Path
 from typing import Self, TypeVar
 
+import jax
+
 from xax.core.conf import field, get_run_dir
 from xax.core.state import State
 from xax.nn.parallel import is_master
@@ -19,6 +21,7 @@ from xax.utils.text import show_info
 logger = logging.getLogger(__name__)
 
 
+@jax.tree_util.register_dataclass
 @dataclass
 class ArtifactsConfig(BaseConfig):
     exp_dir: str | None = field(None, help="The fixed experiment directory")
@@ -43,14 +46,14 @@ class ArtifactsMixin(BaseTask[Config]):
             run_dir = Path(task_file).resolve().parent
         return run_dir / self.task_name
 
-    def set_exp_dir(self, exp_dir: Path) -> Self:
-        self._exp_dir = exp_dir
-        return self
-
     @property
     def exp_dir(self) -> Path:
         return self.get_exp_dir()
 
+    def set_exp_dir(self, exp_dir: str | Path) -> Self:
+        self._exp_dir = Path(exp_dir).expanduser().resolve()
+        return self
+
     def add_lock_file(self, lock_type: str, *, exists_ok: bool = False) -> None:
         if (lock_file := self.exp_dir / f".lock_{lock_type}").exists():
             if not exists_ok: