xax 0.0.7__py3-none-any.whl → 0.1.0__py3-none-any.whl

xax/task/script.py

@@ -17,8 +17,6 @@ from xax.task.mixins import (
     ProcessMixin,
     RunnableConfig,
     RunnableMixin,
-    StepContextConfig,
-    StepContextMixin,
 )
 
 
@@ -28,7 +26,6 @@ class ScriptConfig(
     GPUStatsConfig,
     ProcessConfig,
     LoggerConfig,
-    StepContextConfig,
     ArtifactsConfig,
     RunnableConfig,
     BaseConfig,
@@ -44,7 +41,6 @@ class Script(
     GPUStatsMixin[ConfigT],
     ProcessMixin[ConfigT],
     LoggerMixin[ConfigT],
-    StepContextMixin[ConfigT],
     ArtifactsMixin[ConfigT],
     RunnableMixin[ConfigT],
     BaseTask[ConfigT],

xax/utils/debugging.py

@@ -0,0 +1,49 @@
+"""Defines some useful Jax debugging utilities."""
+
+from collections import deque
+from collections.abc import Iterable, Mapping
+from typing import Any, Callable, Deque
+
+from jaxtyping import Array
+
+
+def get_named_leaves(
+    obj: Any,  # noqa: ANN401
+    is_leaf: Callable[[Any], bool] = lambda x: isinstance(x, Array),  # noqa: ANN401
+    max_depth: int = 100,
+) -> list[tuple[str, Any]]:  # noqa: ANN401
+    ret: list[tuple[str, Any]] = []
+    q: Deque[tuple[int, str, Any]] = deque()  # noqa: ANN401
+    q.append((0, "", obj))
+
+    while q:
+        depth, name, node = q.popleft()
+
+        if depth > max_depth:
+            continue
+
+        if hasattr(node, "__dict__") and isinstance(node.__dict__, Mapping):
+            for cname, cnode in node.__dict__.items():
+                gname = f"{name}.{cname}" if name else cname
+                if is_leaf(cnode):
+                    ret.append((gname, cnode))
+                else:
+                    q.append((depth + 1, gname, cnode))
+
+        elif isinstance(node, Mapping):
+            for cname, cnode in node.items():
+                gname = f"{name}.{cname}" if name else cname
+                if is_leaf(cnode):
+                    ret.append((gname, cnode))
+                else:
+                    q.append((depth + 1, gname, cnode))
+
+        elif isinstance(node, Iterable):
+            for i, cnode in enumerate(node):
+                gname = f"{name}.{i}" if name else str(i)
+                if is_leaf(cnode):
+                    ret.append((gname, cnode))
+                else:
+                    q.append((depth + 1, gname, cnode))
+
+    return ret

xax/utils/experiments.py

@@ -23,7 +23,8 @@ import urllib.request
 import warnings
 from abc import ABC, abstractmethod
 from pathlib import Path
-from typing import Any, Iterator, TypeVar, cast
+from types import TracebackType
+from typing import Any, Iterator, Self, TypeVar, cast
 from urllib.parse import urlparse
 
 import git
@@ -116,19 +117,19 @@ class StateTimer:
         logs: dict[str, dict[str, int | float]] = {}
 
         # Logs step statistics.
-        logs["⏰ steps"] = {
+        logs["⌛ steps"] = {
             "total": self.step_timer.steps,
             "per-second": self.step_timer.steps_per_second,
         }
 
         # Logs sample statistics.
-        logs["⏰ samples"] = {
+        logs["⌛ samples"] = {
             "total": self.sample_timer.steps,
             "per-second": self.sample_timer.steps_per_second,
         }
 
         # Logs full iteration statistics.
-        logs["🔧 dt"] = {
+        logs["⌛ dt"] = {
             "iter": self.iter_timer.iter_seconds,
         }
 
@@ -147,6 +148,24 @@ class IntervalTicker:
         return False
 
 
+class ContextTimer:
+    def __init__(self) -> None:
+        self.start_time = 0.0
+        self.elapsed_time = 0.0
+
+    def __enter__(self) -> Self:
+        self.start_time = time.time()
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_value: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> None:
+        self.elapsed_time = time.time() - self.start_time
+
+
 def abs_path(path: str) -> str:
     return str(Path(path).resolve())
 

xax/utils/jaxpr.py

@@ -0,0 +1,77 @@
+"""Visualize JAXPR."""
+
+from pathlib import Path
+
+import jax
+import jax.core
+
+
+def save_jaxpr_dot(closed_jaxpr: jax.core.ClosedJaxpr, filename: str | Path) -> None:
+    """Save the JAXPR to a DOT file.
+
+    Example usage:
+
+        grad_fn_jaxpr = jax.make_jaxpr(loss_fn)(variables)
+        save_jaxpr_dot(grad_fn_jaxpr, "grad_fn_jaxpr.dot")
+
+    Then, you can visualize the JAXPR using Graphviz:
+
+        dot -Tpng grad_fn_jaxpr.dot > grad_fn_jaxpr.png
+
+    Args:
+        closed_jaxpr: The closed JAXPR to save.
+        filename: The filename to save the JAXPR to.
+    """
+    if hasattr(closed_jaxpr, "jaxpr"):
+        jaxpr = closed_jaxpr.jaxpr
+    else:
+        jaxpr = closed_jaxpr
+
+    with open(filename, "w") as f:
+        f.write("digraph Jaxpr {\n")
+
+        var_names: dict[jax.core.Var, str] = {}
+        var_count = 0
+
+        def get_var_name(var: jax.core.Var) -> str:
+            """Get a unique name for a variable."""
+            nonlocal var_names, var_count
+
+            # Handle Literal objects specially since they're not hashable
+            if isinstance(var, jax.core.Literal):
+                # Create a name based on the literal value
+                name = f"lit_{var.val}"
+                return name
+
+            # For other variables
+            if var not in var_names:
+                name = f"var_{var_count}"
+                var_names[var] = name
+                var_count += 1
+            return var_names[var]
+
+        for var in jaxpr.invars:
+            node_name = get_var_name(var)
+            f.write(f'  {node_name} [label="{node_name}\\n(input)"];\n')
+
+        eq_count = 0
+        for eq in jaxpr.eqns:
+            eq_node = f"eq{eq_count}"
+            label = f"{eq.primitive.name}"
+            f.write(f'  {eq_node} [shape=box, label="{label}"];\n')
+
+            for invar in eq.invars:
+                var_name = get_var_name(invar)
+                f.write(f"  {var_name} -> {eq_node};\n")
+
+            for outvar in eq.outvars:
+                var_name = get_var_name(outvar)
+                f.write(f"  {eq_node} -> {var_name};\n")
+
+            eq_count += 1
+
+        for var in jaxpr.outvars:
+            node_name = get_var_name(var)
+            f.write(f'  {node_name} [peripheries=2, label="{node_name}\\n(output)"];\n')
+
+        f.write("}\n")

xax/utils/pytree.py

@@ -4,7 +4,7 @@ import chex
 import jax
 import jax.numpy as jnp
 from jax import Array
-from jaxtyping import PyTree
+from jaxtyping import PRNGKeyArray, PyTree
 
 
 def slice_array(x: Array, start: Array, slice_length: int) -> Array:
@@ -12,6 +12,14 @@ def slice_array(x: Array, start: Array, slice_length: int) -> Array:
 
     For multi-dimensional arrays, this slices only along the first dimension
     and keeps all other dimensions intact.
+
+    Args:
+        x: The array to slice.
+        start: The start index of the slice.
+        slice_length: The length of the slice.
+
+    Returns:
+        The sliced array.
     """
     chex.assert_shape(start, ())
     chex.assert_shape(slice_length, ())
@@ -38,6 +46,21 @@ def flatten_pytree(pytree: PyTree, flatten_size: int) -> PyTree:
     return jax.tree_util.tree_map(lambda x: flatten_array(x, flatten_size), pytree)
 
 
+def pytree_has_nans(pytree: PyTree) -> Array:
+    """Check if a pytree has any NaNs."""
+    has_nans = jax.tree_util.tree_reduce(
+        lambda a, b: jnp.logical_or(a, b),
+        jax.tree_util.tree_map(lambda x: jnp.any(jnp.isnan(x)), pytree),
+    )
+    return has_nans
+
+
+def update_pytree(cond: Array, new: PyTree, original: PyTree) -> PyTree:
+    """Update a pytree based on a condition."""
+    # Tricky, need use tree_map because where expects array leafs.
+    return jax.tree_util.tree_map(lambda x, y: jnp.where(cond, x, y), new, original)
+
+
 def compute_nan_ratio(pytree: PyTree) -> Array:
     """Computes the ratio of NaNs vs non-NaNs in a given PyTree."""
     nan_counts = jax.tree_util.tree_map(lambda x: jnp.sum(jnp.isnan(x)), pytree)
@@ -48,3 +71,168 @@ def compute_nan_ratio(pytree: PyTree) -> Array:
     overall_nan_ratio = jnp.array(total_nans / total_elements)
 
     return overall_nan_ratio
+
+
+def reshuffle_pytree(data: PyTree, batch_shape: tuple[int, ...], rng: PRNGKeyArray) -> PyTree:
+    """Reshuffle a pytree along the leading dimensions.
+
+    This function reshuffles the data along the leading dimensions specified by batch_shape.
+    Assumes the dimensions to shuffle are the leading ones.
+
+    Args:
+        data: A pytree with arrays.
+        batch_shape: A tuple of integers specifying the size of each leading dimension to shuffle.
+        rng: A JAX PRNG key.
+
+    Returns:
+        A new pytree with the same structure but with the data reshuffled along
+        the leading dimensions.
+    """
+    # Create a permutation for the flattened batch dimensions
+    flat_size = 1
+    for dim in batch_shape:
+        flat_size *= dim
+
+    perm = jax.random.permutation(rng, flat_size)
+
+    def permute_array(x: Array) -> Array:
+        if not isinstance(x, jnp.ndarray):
+            return x
+
+        # Check if the array has enough dimensions
+        if len(x.shape) < len(batch_shape):
+            return x
+
+        # Check if the dimensions match the batch_shape
+        for i, dim in enumerate(batch_shape):
+            if x.shape[i] != dim:
+                return x
+
+        # Reshape to flatten the batch dimensions
+        orig_shape = x.shape
+        reshaped = x.reshape((flat_size,) + orig_shape[len(batch_shape) :])
+
+        # Apply the permutation
+        permuted = reshaped[perm]
+
+        # Reshape back to the original shape
+        return permuted.reshape(orig_shape)
+
+    return jax.tree_util.tree_map(permute_array, data)
+
+
+def reshuffle_pytree_independently(data: PyTree, batch_shape: tuple[int, ...], rng: PRNGKeyArray) -> PyTree:
+    """Reshuffle a rollout array across arbitrary batch dimensions independently of each other."""
+    rngs = jax.random.split(rng, len(batch_shape))
+    perms = [jax.random.permutation(rng_i, dim) for rng_i, dim in zip(rngs, batch_shape)]
+    # n-dimensional index grid from permutations
+    idx_grids = jnp.meshgrid(*perms, indexing="ij")
+
+    def permute_array(x: Array) -> Array:
+        if isinstance(x, Array):
+            return x[tuple(idx_grids)]
+        return x
+
+    return jax.tree_util.tree_map(permute_array, data)
+
+
+TransposeResult = tuple[PyTree, tuple[int, ...], tuple[int, ...]]
+PathType = tuple[str | int, ...]
+
+
+def reshuffle_pytree_along_dims(
+    data: PyTree,
+    dims: tuple[int, ...],
+    shape_dims: tuple[int, ...],
+    rng: PRNGKeyArray,
+) -> PyTree:
+    """Reshuffle a pytree along arbitrary dimensions.
+
+    Allows reshuffling along any dimensions, not just the leading ones.
+    It transposes the data to make the specified dimensions the leading ones,
+    then reshuffles, and finally transposes back.
+
+    Args:
+        data: A pytree with arrays.
+        dims: A tuple of integers specifying which dimensions to shuffle along.
+            For example, (1,) would shuffle along the second dimension.
+        shape_dims: A tuple of integers specifying the size of each dimension to shuffle.
+            Must have the same length as dims.
+        rng: A JAX PRNG key.
+
+    Returns:
+        A new pytree with the same structure but with the data reshuffled along
+        the specified dimensions.
+    """
+    if len(dims) != len(shape_dims):
+        raise ValueError(f"dims {dims} and shape_dims {shape_dims} must have the same length")
+
+    def transpose_for_shuffle(x: PyTree) -> TransposeResult:
+        if not isinstance(x, jnp.ndarray):
+            return x, (), ()
+
+        # Check if the array has enough dimensions
+        if len(x.shape) <= max(dims):
+            return x, (), ()
+
+        # Check if the dimensions match the shape_dims
+        for i, dim in enumerate(dims):
+            if x.shape[dim] != shape_dims[i]:
+                raise ValueError(f"Array shape {x.shape} doesn't match shape_dims {shape_dims} at dimension {dim}")
+
+        # Create the transpose order to move the specified dimensions to the front
+        # while preserving the relative order of the other dimensions
+        n_dims = len(x.shape)
+        other_dims = [i for i in range(n_dims) if i not in dims]
+        transpose_order = tuple(dims) + tuple(other_dims)
+
+        # Transpose the array
+        transposed = jnp.transpose(x, transpose_order)
+
+        return transposed, transpose_order, x.shape
+
+    def transpose_back(x: PyTree, transpose_order: tuple[int, ...], original_shape: tuple[int, ...]) -> PyTree:
+        if not isinstance(x, jnp.ndarray) or not transpose_order:
+            return x
+
+        # Create the inverse transpose order
+        inverse_order = [0] * len(transpose_order)
+        for i, j in enumerate(transpose_order):
+            inverse_order[j] = i
+
+        # Transpose back
+        return jnp.transpose(x, inverse_order)
+
+    # First, transpose all arrays to make the specified dimensions the leading ones
+    transposed_data: dict[PathType, Array] = {}
+    transpose_info: dict[PathType, tuple[tuple[int, ...], tuple[int, ...]]] = {}
+
+    def prepare_for_shuffle(path: PathType, x: PyTree) -> PyTree:
+        if isinstance(x, jnp.ndarray):
+            transposed, transpose_order, original_shape = transpose_for_shuffle(x)
+            if isinstance(transposed, jnp.ndarray):  # Check if it's an array
+                transposed_data[path] = transposed
+                transpose_info[path] = (transpose_order, original_shape)
+        return x
+
+    jax.tree_util.tree_map_with_path(prepare_for_shuffle, data)
+
+    # Create a transposed pytree
+    def get_transposed(path: PathType, x: PyTree) -> PyTree:
+        if isinstance(x, jnp.ndarray) and path in transposed_data:
+            return transposed_data[path]
+        return x
+
+    transposed_pytree = jax.tree_util.tree_map_with_path(get_transposed, data)
+
+    # Reshuffle the transposed pytree along the leading dimensions
+    reshuffled_transposed = reshuffle_pytree(transposed_pytree, shape_dims, rng)
+
+    # Transpose back
+    def restore_transpose(path: PathType, x: PyTree) -> PyTree:
+        if isinstance(x, jnp.ndarray) and path in transpose_info:
+            transpose_order, original_shape = transpose_info[path]
+            return transpose_back(x, transpose_order, original_shape)
+        return x
+
+    return jax.tree_util.tree_map_with_path(restore_transpose, reshuffled_transposed)

xax/utils/tensorboard.py

@@ -14,7 +14,7 @@ from PIL.Image import Image as PILImage
 from tensorboard.compat.proto.config_pb2 import RunMetadata
 from tensorboard.compat.proto.event_pb2 import Event, TaggedRunMetadata
 from tensorboard.compat.proto.graph_pb2 import GraphDef
-from tensorboard.compat.proto.summary_pb2 import Summary, SummaryMetadata
+from tensorboard.compat.proto.summary_pb2 import HistogramProto, Summary, SummaryMetadata
 from tensorboard.compat.proto.tensor_pb2 import TensorProto
 from tensorboard.compat.proto.tensor_shape_pb2 import TensorShapeProto
 from tensorboard.plugins.text.plugin_data_pb2 import TextPluginData
@@ -25,6 +25,65 @@ from xax.core.state import Phase
 ImageShape = Literal["HWC", "CHW", "HW", "NHWC", "NCHW", "NHW"]
 
 
+def make_histogram(values: np.ndarray, bins: str | np.ndarray, max_bins: int | None = None) -> HistogramProto:
+    """Convert values into a histogram proto using logic from histogram.cc.
+
+    Args:
+        values: Input values to create histogram from
+        bins: Bin specification (string like 'auto' or array of bin edges)
+        max_bins: Optional maximum number of bins
+
+    Returns:
+        HistogramProto containing the histogram data
+    """
+    if values.size == 0:
+        raise ValueError("The input has no element.")
+    values = values.reshape(-1)
+    counts, limits = np.histogram(values, bins=bins)
+    num_bins = len(counts)
+
+    if max_bins is not None and num_bins > max_bins:
+        subsampling = num_bins // max_bins
+        subsampling_remainder = num_bins % subsampling
+        if subsampling_remainder != 0:
+            counts = np.pad(
+                counts,
+                pad_width=[[0, subsampling - subsampling_remainder]],
+                mode="constant",
+                constant_values=0,
+            )
+        counts = counts.reshape(-1, subsampling).sum(axis=-1)
+        new_limits = np.empty((counts.size + 1,), limits.dtype)
+        new_limits[:-1] = limits[:-1:subsampling]
+        new_limits[-1] = limits[-1]
+        limits = new_limits
+
+    # Find the first and the last bin defining the support of the histogram:
+    cum_counts = np.cumsum(np.greater(counts, 0))
+    start, end = np.searchsorted(cum_counts, [0, cum_counts[-1] - 1], side="right")
+    start = int(start)
+    end = int(end) + 1
+
+    # TensorBoard only includes the right bin limits. To still have the leftmost limit
+    # included, we include an empty bin left.
+    counts = counts[start - 1 : end] if start > 0 else np.concatenate([[0], counts[:end]])
+    limits = limits[start : end + 1]
+
+    if counts.size == 0 or limits.size == 0:
+        raise ValueError("The histogram is empty, please file a bug report.")
+
+    sum_sq = values.dot(values)
+    return HistogramProto(
+        min=values.min(),
+        max=values.max(),
+        num=len(values),
+        sum=values.sum(),
+        sum_squares=sum_sq,
+        bucket_limit=limits.tolist(),
+        bucket=counts.tolist(),
+    )
+
+
 class TensorboardProtobufWriter:
     def __init__(
         self,
@@ -263,6 +322,123 @@ class TensorboardWriter:
             walltime=walltime,
         )
 
+    def add_histogram(
+        self,
+        tag: str,
+        values: np.ndarray,
+        global_step: int | None = None,
+        bins: str | np.ndarray = "auto",
+        walltime: float | None = None,
+        max_bins: int | None = None,
+    ) -> None:
+        hist = make_histogram(values.astype(float), bins, max_bins)
+        self.pb_writer.add_summary(
+            Summary(value=[Summary.Value(tag=tag, histo=hist)]),
+            global_step=global_step,
+            walltime=walltime,
+        )
+
+    def add_histogram_raw(
+        self,
+        tag: str,
+        min: float | np.ndarray,
+        max: float | np.ndarray,
+        num: int | np.ndarray,
+        sum: float | np.ndarray,
+        sum_squares: float | np.ndarray,
+        bucket_limits: list[float | np.ndarray],
+        bucket_counts: list[int | np.ndarray],
+        global_step: int | None = None,
+        walltime: float | None = None,
+    ) -> None:
+        """Add histogram with raw data to summary.
+
+        Args:
+            tag: Data identifier
+            min: Min value
+            max: Max value
+            num: Number of values
+            sum: Sum of all values
+            sum_squares: Sum of squares for all values
+            bucket_limits: Upper value per bucket
+            bucket_counts: Number of values per bucket
+            global_step: Global step value to record
+            walltime: Optional override default walltime
+        """
+        if len(bucket_limits) != len(bucket_counts):
+            raise ValueError("len(bucket_limits) != len(bucket_counts)")
+
+        # Convert numpy arrays to Python types
+        hist = HistogramProto(
+            min=float(min),
+            max=float(max),
+            num=int(num),
+            sum=float(sum),
+            sum_squares=float(sum_squares),
+            bucket_limit=[float(x) for x in bucket_limits],
+            bucket=[int(x) for x in bucket_counts],
+        )
+        self.pb_writer.add_summary(
+            Summary(value=[Summary.Value(tag=tag, histo=hist)]),
+            global_step=global_step,
+            walltime=walltime,
+        )
+
+    def add_gaussian_distribution(
+        self,
+        tag: str,
+        mean: float | np.ndarray,
+        std: float | np.ndarray,
+        bins: int = 16,
+        global_step: int | None = None,
+        walltime: float | None = None,
+    ) -> None:
+        """Add a Gaussian distribution to the summary.
+
+        Args:
+            tag: Data identifier
+            mean: Mean of the Gaussian distribution
+            std: Standard deviation of the Gaussian distribution
+            bins: Number of bins to use for the histogram
+            global_step: Global step value to record
+            walltime: Optional override default walltime
+        """
+        # Convert numpy arrays to Python types
+        mean = float(mean)
+        std = float(std)
+
+        # Create bin edges spanning ±3 standard deviations
+        bin_edges = np.linspace(mean - 3 * std, mean + 3 * std, bins + 1)
+
+        # Calculate the probability density for each bin
+        bin_centers = (bin_edges[1:] + bin_edges[:-1]) / 2
+        gaussian_pdf = np.exp(-0.5 * ((bin_centers - mean) / std) ** 2) / (std * np.sqrt(2 * np.pi))
+
+        # Scale the PDF to represent counts
+        num_samples = bins * 1000
+        bucket_counts = (gaussian_pdf * num_samples * (bin_edges[1] - bin_edges[0])).astype(int)
+
+        # Ensure we have at least one count per bin for visualization
+        bucket_counts = np.maximum(bucket_counts, 1)
+
+        # Calculate actual statistics based on the discretized distribution
+        total_counts = float(bucket_counts.sum())
+        weighted_sum = float((bin_centers * bucket_counts).sum())
+        weighted_sum_squares = float((bin_centers**2 * bucket_counts).sum())
+
+        self.add_histogram_raw(
+            tag=tag,
+            min=float(bin_edges[0]),
+            max=float(bin_edges[-1]),
+            num=int(total_counts),
+            sum=weighted_sum,
+            sum_squares=weighted_sum_squares,
+            bucket_limits=bin_edges[1:].tolist(),  # TensorBoard expects right bin edges
+            bucket_counts=bucket_counts.tolist(),
+            global_step=global_step,
+            walltime=walltime,
+        )
+
 
 class TensorboardWriterKwargs(TypedDict):
     max_queue_size: int

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

@@ -1,12 +1,13 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.4
 Name: xax
-Version: 0.0.7
-Summary: The xax project
-Home-page: https://github.com/dpshai/xax
+Version: 0.1.0
+Summary: A library for fast Jax experimentation
+Home-page: https://github.com/kscalelabs/xax
 Author: Benjamin Bolte
 Requires-Python: >=3.11
 Description-Content-Type: text/markdown
 License-File: LICENSE
+Requires-Dist: attrs
 Requires-Dist: jax
 Requires-Dist: jaxtyping
 Requires-Dist: equinox
@@ -30,10 +31,28 @@ Requires-Dist: pytest; extra == "dev"
 Requires-Dist: types-pillow; extra == "dev"
 Requires-Dist: types-psutil; extra == "dev"
 Requires-Dist: types-requests; extra == "dev"
+Provides-Extra: export
+Requires-Dist: orbax-export; extra == "export"
+Requires-Dist: tensorflow; extra == "export"
+Provides-Extra: flax
+Requires-Dist: flax; extra == "flax"
+Provides-Extra: all
+Requires-Dist: black; extra == "all"
+Requires-Dist: darglint; extra == "all"
+Requires-Dist: mypy; extra == "all"
+Requires-Dist: ruff; extra == "all"
+Requires-Dist: pytest; extra == "all"
+Requires-Dist: types-pillow; extra == "all"
+Requires-Dist: types-psutil; extra == "all"
+Requires-Dist: types-requests; extra == "all"
+Requires-Dist: orbax-export; extra == "all"
+Requires-Dist: tensorflow; extra == "all"
+Requires-Dist: flax; extra == "all"
 Dynamic: author
 Dynamic: description
 Dynamic: description-content-type
 Dynamic: home-page
+Dynamic: license-file
 Dynamic: provides-extra
 Dynamic: requires-dist
 Dynamic: requires-python