xax 0.1.15__tar.gz → 0.2.0__tar.gz

{xax-0.1.15/xax.egg-info → xax-0.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: xax
-Version: 0.1.15
+Version: 0.2.0
 Summary: A library for fast Jax experimentation
 Home-page: https://github.com/kscalelabs/xax
 Author: Benjamin Bolte
@@ -8,14 +8,14 @@ Requires-Python: >=3.11
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: attrs
+Requires-Dist: chex
+Requires-Dist: dpshdl
+Requires-Dist: equinox
+Requires-Dist: importlib-resources
 Requires-Dist: jax
 Requires-Dist: jaxtyping
-Requires-Dist: equinox
 Requires-Dist: optax
-Requires-Dist: dpshdl
-Requires-Dist: chex
-Requires-Dist: importlib-resources
-Requires-Dist: cloudpickle
+Requires-Dist: orbax-checkpoint
 Requires-Dist: pillow
 Requires-Dist: omegaconf
 Requires-Dist: gitpython

{xax-0.1.15 → xax-0.2.0}/pyproject.toml

@@ -35,7 +35,6 @@ explicit_package_bases = true
 [[tool.mypy.overrides]]
 
 module = [
-    "cloudpickle.*",
     "optax.*",
     "setuptools.*",
     "tensorboard.*",

{xax-0.1.15 → xax-0.2.0}/xax/__init__.py

@@ -12,7 +12,7 @@ and running the update script:
     python -m scripts.update_api --inplace
 """
 
-__version__ = "0.1.15"
+__version__ = "0.2.0"
 
 # This list shouldn't be modified by hand; instead, run the update script.
 __all__ = [

{xax-0.1.15 → xax-0.2.0}/xax/core/state.py

@@ -12,6 +12,14 @@ from xax.core.conf import field
 Phase = Literal["train", "valid"]
 
 
+def _phase_to_int(phase: Phase) -> int:
+    return {"train": 0, "valid": 1}[phase]
+
+
+def _int_to_phase(i: int) -> Phase:
+    return cast(Phase, ["train", "valid"][i])
+
+
 class StateDict(TypedDict, total=False):
     num_steps: NotRequired[int]
     num_samples: NotRequired[int]
@@ -35,7 +43,7 @@ class State:
 
     @property
     def phase(self) -> Phase:
-        return cast(Phase, ["train", "valid"][self._phase])
+        return _int_to_phase(self._phase)
 
     @classmethod
     def init_state(cls) -> "State":
@@ -74,3 +82,20 @@ class State:
                 case _:
                     raise ValueError(f"Invalid phase: {phase}")
         return State(**{**asdict(self), **kwargs, **extra_kwargs})
+
+    def to_dict(self) -> dict[str, int | float | str]:
+        return {
+            "num_steps": int(self.num_steps),
+            "num_samples": int(self.num_samples),
+            "num_valid_steps": int(self.num_valid_steps),
+            "num_valid_samples": int(self.num_valid_samples),
+            "start_time_s": float(self.start_time_s),
+            "elapsed_time_s": float(self.elapsed_time_s),
+            "phase": str(self.phase),
+        }
+
+    @classmethod
+    def from_dict(cls, d: dict[str, int | float | str]) -> "State":
+        if "phase" in d:
+            d["_phase"] = _phase_to_int(cast(Phase, d.pop("phase")))
+        return cls(**d)  # type: ignore[arg-type]

{xax-0.1.15 → xax-0.2.0}/xax/requirements.txt

@@ -2,16 +2,16 @@
 
 # Core ML/JAX dependencies
 attrs
+chex
+dpshdl
+equinox
+importlib-resources
 jax
 jaxtyping
-equinox
 optax
-dpshdl
-chex
-importlib-resources
+orbax-checkpoint
 
 # Data processing and serialization
-cloudpickle
 pillow
 
 # Configuration and project management

{xax-0.1.15 → xax-0.2.0}/xax/task/base.py

@@ -79,7 +79,7 @@ class BaseTask(Generic[Config]):
     def on_training_end(self, state: State) -> State:
         return state
 
-    def on_after_checkpoint_save(self, ckpt_path: Path, state: State) -> State:
+    def on_after_checkpoint_save(self, ckpt_path: Path, state: State | None) -> State | None:
         return state
 
     @functools.cached_property

{xax-0.1.15 → xax-0.2.0}/xax/task/logger.py

@@ -18,11 +18,22 @@ 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, cast, get_args
+from typing import (
+    Any,
+    Callable,
+    Iterator,
+    Literal,
+    Self,
+    Sequence,
+    TypeVar,
+    cast,
+    get_args,
+)
 
 import jax
 import jax.numpy as jnp
 import numpy as np
+from jax._src.core import ClosedJaxpr
 from jaxtyping import Array
 from PIL import Image, ImageDraw, ImageFont
 from PIL.Image import Image as PILImage
@@ -194,7 +205,10 @@ def tile_images(images: list[PILImage], sep: int = 0) -> PILImage:
     return tiled
 
 
-def as_numpy(array: Array) -> np.ndarray:
+def as_numpy(array: Array | np.ndarray) -> np.ndarray:
+    """Convert a JAX array or numpy array to numpy array."""
+    if isinstance(array, np.ndarray):
+        return array
     array = jax.device_get(array)
     if jax.dtypes.issubdtype(array.dtype, jnp.floating):
         array = array.astype(jnp.float32)
@@ -205,6 +219,19 @@ def as_numpy(array: Array) -> np.ndarray:
     return np.array(array)
 
 
+def as_numpy_opt(array: Array | np.ndarray | None) -> np.ndarray | None:
+    """Convert an optional JAX array or numpy array to numpy array."""
+    if array is None:
+        return None
+    return as_numpy(array)
+
+
+@dataclass(kw_only=True)
+class LogString:
+    value: str
+    secondary: bool
+
+
 @dataclass(kw_only=True)
 class LogImage:
     image: PILImage
@@ -223,6 +250,12 @@ class LogVideo:
     fps: int
 
 
+@dataclass(kw_only=True)
+class LogScalar:
+    value: Number
+    secondary: bool
+
+
 @dataclass(kw_only=True)
 class LogDistribution:
     mean: Number
@@ -240,15 +273,29 @@ class LogHistogram:
     bucket_counts: list[int]
 
 
+@dataclass(kw_only=True)
+class LogMesh:
+    vertices: np.ndarray
+    colors: np.ndarray | None
+    faces: np.ndarray | None
+    config_dict: dict[str, Any] | None  # noqa: ANN401
+
+
+@dataclass(kw_only=True)
+class LogGraph:
+    computation: ClosedJaxpr
+
+
 @dataclass(kw_only=True)
 class LogLine:
     state: State
-    scalars: dict[str, dict[str, Number]]
+    scalars: dict[str, dict[str, LogScalar]]
     distributions: dict[str, dict[str, LogDistribution]]
     histograms: dict[str, dict[str, LogHistogram]]
-    strings: dict[str, dict[str, str]]
+    strings: dict[str, dict[str, LogString]]
     images: dict[str, dict[str, LogImage]]
     videos: dict[str, dict[str, LogVideo]]
+    meshes: dict[str, dict[str, LogMesh]]
 
 
 @dataclass(kw_only=True)
@@ -515,12 +562,13 @@ class Logger:
     """Defines an intermediate container which holds values to log somewhere else."""
 
     def __init__(self, default_namespace: str = DEFAULT_NAMESPACE) -> None:
-        self.scalars: dict[str, dict[str, Callable[[], Number]]] = defaultdict(dict)
+        self.scalars: dict[str, dict[str, Callable[[], LogScalar]]] = defaultdict(dict)
         self.distributions: dict[str, dict[str, Callable[[], LogDistribution]]] = defaultdict(dict)
         self.histograms: dict[str, dict[str, Callable[[], LogHistogram]]] = defaultdict(dict)
-        self.strings: dict[str, dict[str, Callable[[], str]]] = defaultdict(dict)
+        self.strings: dict[str, dict[str, Callable[[], LogString]]] = defaultdict(dict)
         self.images: dict[str, dict[str, Callable[[], LogImage]]] = defaultdict(dict)
         self.videos: dict[str, dict[str, Callable[[], LogVideo]]] = defaultdict(dict)
+        self.meshes: dict[str, dict[str, Callable[[], LogMesh]]] = defaultdict(dict)
         self.default_namespace = default_namespace
         self.loggers: list[LoggerImpl] = []
 
@@ -548,6 +596,7 @@ class Logger:
             strings={k: {kk: v() for kk, v in v.items()} for k, v in self.strings.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()},
+            meshes={k: {kk: v() for kk, v in v.items()} for k, v in self.meshes.items()},
         )
 
     def clear(self) -> None:
@@ -557,6 +606,7 @@ class Logger:
         self.strings.clear()
         self.images.clear()
         self.videos.clear()
+        self.meshes.clear()
 
     def write(self, state: State) -> None:
         """Writes the current step's logging information.
@@ -616,13 +666,23 @@ class Logger:
     def resolve_namespace(self, namespace: str | None = None) -> str:
         return "_".join([self.default_namespace if namespace is None else namespace] + NAMESPACE_STACK)
 
-    def log_scalar(self, key: str, value: Callable[[], Number] | Number, *, namespace: str | None = None) -> None:
+    def log_scalar(
+        self,
+        key: str,
+        value: Callable[[], Number] | Number,
+        *,
+        namespace: str | None = None,
+        secondary: bool = False,
+    ) -> None:
         """Logs a scalar value.
 
         Args:
             key: The key being logged
             value: The scalar value being logged
             namespace: An optional logging namespace
+            secondary: If set, treat this as a secondary value (meaning, it is
+                less important than other values, and some downstream loggers
+                will not display it)
         """
         if not self.active:
             raise RuntimeError("The logger is not active")
@@ -632,11 +692,11 @@ class Logger:
             assert value.ndim == 0, f"Scalar must be a 0D array, got shape {value.shape}"
 
         @functools.lru_cache(maxsize=None)
-        def scalar_future() -> Number:
+        def scalar_future() -> LogScalar:
             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
+            return LogScalar(value=value_concrete, secondary=secondary)
 
         self.scalars[namespace][key] = scalar_future
 
@@ -770,21 +830,31 @@ class Logger:
 
         self.histograms[namespace][key] = histogram_future
 
-    def log_string(self, key: str, value: Callable[[], str] | str, *, namespace: str | None = None) -> None:
+    def log_string(
+        self,
+        key: str,
+        value: Callable[[], str] | str,
+        *,
+        namespace: str | None = None,
+        secondary: bool = False,
+    ) -> None:
         """Logs a string value.
 
         Args:
             key: The key being logged
             value: The string value being logged
             namespace: An optional logging namespace
+            secondary: If set, treat this as a secondary value (meaning, it is
+                less important than other values, and some downstream loggers
+                will not display it)
         """
         if not self.active:
             raise RuntimeError("The logger is not active")
         namespace = self.resolve_namespace(namespace)
 
         @functools.lru_cache(maxsize=None)
-        def value_future() -> str:
-            return value() if callable(value) else value
+        def value_future() -> LogString:
+            return LogString(value=value() if callable(value) else value, secondary=secondary)
 
         self.strings[namespace][key] = value_future
 
@@ -1019,6 +1089,73 @@ class Logger:
 
         self.videos[namespace][key] = video_future
 
+    def log_mesh(
+        self,
+        key: str,
+        vertices: np.ndarray | Array | Callable[[], np.ndarray | Array],
+        colors: np.ndarray | Array | None | Callable[[], np.ndarray | Array | None] = None,
+        faces: np.ndarray | Array | None | Callable[[], np.ndarray | Array | None] = None,
+        config_dict: dict[str, Any] | None = None,
+        *,
+        namespace: str | None = None,
+    ) -> None:
+        if not self.active:
+            raise RuntimeError("The logger is not active")
+        namespace = self.resolve_namespace(namespace)
+
+        @functools.lru_cache(maxsize=None)
+        def mesh_future() -> LogMesh:
+            with ContextTimer() as timer:
+                # Get the raw values
+                vertices_val = vertices() if callable(vertices) else vertices
+                colors_val = colors() if callable(colors) else colors
+                faces_val = faces() if callable(faces) else faces
+
+                # Convert to numpy arrays with proper type handling
+                vertices_np = as_numpy(vertices_val)
+                colors_np = as_numpy_opt(colors_val)
+                faces_np = as_numpy_opt(faces_val)
+
+                # Checks vertices shape.
+                if vertices_np.ndim == 2:
+                    vertices_np = vertices_np[None]
+                if vertices_np.shape[-1] != 3 or vertices_np.ndim != 3:
+                    raise ValueError("Vertices must have shape (N, 3) or (B, N, 3)")
+
+                # Checks colors shape.
+                if colors_np is not None:
+                    if colors_np.ndim == 2:
+                        colors_np = colors_np[None]
+                    if colors_np.shape[-1] != 3 or colors_np.ndim != 3:
+                        raise ValueError("Colors must have shape (N, 3) or (B, N, 3)")
+
+                # Checks faces shape.
+                if faces_np is not None:
+                    if faces_np.ndim == 2:
+                        faces_np = faces_np[None]
+                    if faces_np.shape[-1] != 3 or faces_np.ndim != 3:
+                        raise ValueError("Faces must have shape (N, 3) or (B, N, 3)")
+
+                # Ensures colors dtype is uint8.
+                if colors_np is not None:
+                    if colors_np.dtype != np.uint8:
+                        colors_np = (colors_np * 255).astype(np.uint8)
+
+                # Ensures faces dtype is int32.
+                if faces_np is not None:
+                    if faces_np.dtype != np.int32:
+                        faces_np = faces_np.astype(np.int32)
+
+            logger.debug("Mesh Key: %s, Time: %s", key, timer.elapsed_time)
+            return LogMesh(
+                vertices=vertices_np,
+                colors=colors_np,
+                faces=faces_np,
+                config_dict=config_dict,
+            )
+
+        self.meshes[namespace][key] = mesh_future
+
     def __enter__(self) -> Self:
         self.active = True
         for logger in self.loggers:

{xax-0.1.15 → xax-0.2.0}/xax/task/loggers/json.py

@@ -3,11 +3,19 @@
 import json
 import sys
 from dataclasses import asdict
-from typing import Any, Literal, TextIO
+from typing import Any, Literal, Mapping, TextIO
 
 from jaxtyping import Array
 
-from xax.task.logger import LogError, LoggerImpl, LogLine, LogPing, LogStatus
+from xax.task.logger import (
+    LogError,
+    LoggerImpl,
+    LogLine,
+    LogPing,
+    LogScalar,
+    LogStatus,
+    LogString,
+)
 
 
 def get_json_value(value: Any) -> Any:  # noqa: ANN401
@@ -61,14 +69,14 @@ class JsonLogger(LoggerImpl):
     def get_json(self, line: LogLine) -> str:
         data: dict = {"state": asdict(line.state)}
 
-        def add_logs(log: dict[str, dict[str, Any]], data: dict) -> None:
+        def add_logs(log: Mapping[str, Mapping[str, LogScalar | LogString]], data: dict) -> None:
             for namespace, values in log.items():
                 if self.remove_unicode_from_namespaces:
                     namespace = namespace.encode("ascii", errors="ignore").decode("ascii").strip()
                 if namespace not in data:
                     data[namespace] = {}
                 for k, v in values.items():
-                    data[namespace][k] = get_json_value(v)
+                    data[namespace][k] = get_json_value(v.value)
 
         add_logs(line.scalars, data)
         add_logs(line.strings, data)

{xax-0.1.15 → xax-0.2.0}/xax/task/loggers/stdout.py

@@ -4,11 +4,20 @@ import datetime
 import logging
 import sys
 from collections import deque
-from typing import Any, Deque, TextIO
+from typing import Any, Deque, Mapping, TextIO
 
 from jaxtyping import Array
 
-from xax.task.logger import LogError, LogErrorSummary, LoggerImpl, LogLine, LogPing, LogStatus
+from xax.task.logger import (
+    LogError,
+    LogErrorSummary,
+    LoggerImpl,
+    LogLine,
+    LogPing,
+    LogScalar,
+    LogStatus,
+    LogString,
+)
 from xax.utils.text import Color, colored, format_timedelta
 
 
@@ -95,20 +104,17 @@ class StdoutLogger(LoggerImpl):
     def write_log_window(self, line: LogLine) -> None:
         namespace_to_lines: dict[str, dict[str, str]] = {}
 
-        def add_logs(log: dict[str, dict[str, Any]], namespace_to_lines: dict[str, dict[str, str]]) -> None:
+        def add_logs(
+            log: Mapping[str, Mapping[str, LogScalar | LogString]],
+            namespace_to_lines: dict[str, dict[str, str]],
+        ) -> None:
             for namespace, values in log.items():
-                if not self.log_timers and namespace.startswith("⌛"):
-                    continue
-                if not self.log_perf and namespace.startswith("🔧"):
-                    continue
-                if not self.log_optim and namespace.startswith("📉"):
-                    continue
-                if not self.log_fp and namespace.startswith("⚖️"):
-                    continue
-                if namespace not in namespace_to_lines:
-                    namespace_to_lines[namespace] = {}
                 for k, v in values.items():
-                    v_str = as_str(v, self.precision)
+                    if v.secondary:
+                        continue
+                    if namespace not in namespace_to_lines:
+                        namespace_to_lines[namespace] = {}
+                    v_str = as_str(v.value, self.precision)
                     namespace_to_lines[namespace][k] = v_str
 
         add_logs(line.scalars, namespace_to_lines)
@@ -116,9 +122,8 @@ class StdoutLogger(LoggerImpl):
         if not namespace_to_lines:
             return
 
-        self.write_fp.write("\n")
         for namespace, lines in sorted(namespace_to_lines.items()):
-            self.write_fp.write(f"{colored(namespace, 'cyan', bold=True)}\n")
+            self.write_fp.write(f"\n{colored(namespace, 'cyan', bold=True)}\n")
             for k, v in lines.items():
                 self.write_fp.write(f" ↪ {k}: {v}\n")
 

{xax-0.1.15 → xax-0.2.0}/xax/task/loggers/tensorboard.py

@@ -70,6 +70,9 @@ class TensorboardLogger(LoggerImpl):
         self._started = True
 
     def worker_thread(self) -> None:
+        if os.environ.get("DISABLE_TENSORBOARD", "0") == "1":
+            return
+
         time.sleep(self.wait_seconds)
 
         port = int(os.environ.get("TENSORBOARD_PORT", DEFAULT_TENSORBOARD_PORT))
@@ -158,7 +161,7 @@ class TensorboardLogger(LoggerImpl):
             for scalar_key, scalar_value in scalars.items():
                 writer.add_scalar(
                     f"{namespace}/{scalar_key}",
-                    as_float(scalar_value),
+                    as_float(scalar_value.value),
                     global_step=line.state.num_steps,
                     walltime=walltime,
                 )
@@ -192,7 +195,7 @@ class TensorboardLogger(LoggerImpl):
             for string_key, string_value in strings.items():
                 writer.add_text(
                     f"{namespace}/{string_key}",
-                    string_value,
+                    string_value.value,
                     global_step=line.state.num_steps,
                     walltime=walltime,
                 )
@@ -213,6 +216,19 @@ class TensorboardLogger(LoggerImpl):
                     video_value.frames,
                     fps=video_value.fps,
                     global_step=line.state.num_steps,
+                    walltime=walltime,
+                )
+
+        for namespace, meshes in line.meshes.items():
+            for mesh_key, mesh_value in meshes.items():
+                writer.add_mesh(
+                    f"{namespace}/{mesh_key}",
+                    vertices=mesh_value.vertices,
+                    faces=mesh_value.faces,
+                    colors=mesh_value.colors,
+                    config_dict=mesh_value.config_dict,
+                    global_step=line.state.num_steps,
+                    walltime=walltime,
                 )
 
         for name, contents in self.files.items():