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

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.1.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.1
+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

{xax-0.0.7.dist-info → xax-0.1.1.dist-info}/RECORD

@@ -1,19 +1,22 @@
-xax/__init__.py,sha256=ScTkvKaxgpuKhhs9RINJa2XWCj899ndSYrB3FtScfxw,10509
+xax/__init__.py,sha256=JyKRACir9b0bkuG93bwxADFrVr-Lo76kenDBJtvb_wQ,13280
 xax/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/requirements-dev.txt,sha256=qkscNkFzWd1S5fump-AKH53rR65v2x5FmboFdy_kKvs,128
-xax/requirements.txt,sha256=NmU9PNJhfLtNqqtWWf8WqMjgbBPCn_yt8oMGAgS7Fno,291
+xax/requirements.txt,sha256=9LAEZ5c5gqRSARRVA6xJsVTa4MebPZuC4yOkkwkZJFw,297
 xax/core/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/core/conf.py,sha256=Wuo5WLRWuRTgb8eaihvnG_NZskTu0-P3JkIcl_hKINM,5124
 xax/core/state.py,sha256=y123fL7pMgk25TPG6KN0LRIF_eYnD9eP7OfqtoQJGNE,2178
 xax/nn/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/nn/embeddings.py,sha256=bQGxBFxkLwi2MQLkRfGaHPH5P_KKB21HdI7VNWTKIOQ,11847
+xax/nn/equinox.py,sha256=1Ck6ycz76dhit2LHX4y2lp3WJSPsDuRt7TK7AxxQhww,4837
+xax/nn/export.py,sha256=Do2bLjJTD744mxpQuPYpz8fZ3EIjBLaaZfhp8maNVrg,5303
 xax/nn/functions.py,sha256=CI_OmspaQwN9nl4hwefIU3_I7m6gBZwJ9aGK1JGUgr0,2713
-xax/nn/geom.py,sha256=MtVar9AdqrJQGIFxcIFHyFnV_fblf9Pc4kQT_gTQASI,2195
+xax/nn/geom.py,sha256=eK7I8fUHBc3FT7zpm5Yf__bXFQ4LtX6sa17-DxojLTo,3202
+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=LHDmM2c_Ps5cGEzn_QUpmyInD7zJJm3Yt9eSeij2Vus,7297
-xax/task/logger.py,sha256=orN1jmM4SIR2EiYk8bNoJZscmhX1FytADBU6p9qpows,29256
-xax/task/script.py,sha256=4LyXrpj0V36TjAZT4lvQeiOTqa5U2tommHKwgWDCE24,1025
+xax/task/base.py,sha256=MlH5dTKAiMzFRI5fmXCvL1k8ELbalWMBICeVxmW6k2U,7479
+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
 xax/task/launchers/base.py,sha256=8LB_r6YISKu1vq1zk3aVYmiedRr9MxE5IMRocs6unFI,731
@@ -23,33 +26,35 @@ 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=nxQXkS9JUR38RKsU9qj0dgePKguK0BFa9nl_BdGO8cE,6758
-xax/task/loggers/tensorboard.py,sha256=FGW96z77oG0Kf3cO6Zznx5U3kJNzPWcuSkpY4RnbFCo,6909
+xax/task/loggers/stdout.py,sha256=bR0k-PfmFgLfPxLPb4hZw_8G_msA32UeHfAAu11nEYs,6757
+xax/task/loggers/tensorboard.py,sha256=kI8LvBuBBhPgkP8TeaTQb9SQ0FqaIodwQh2SuWDCnIA,7706
 xax/task/mixins/__init__.py,sha256=D3oU31rB9FeOr9MPLleLt5JFbftUr4sBTwgnwQdc2qA,809
-xax/task/mixins/artifacts.py,sha256=1H7ZbR-KSsXhVtqGVlqMi-TXfn1-dM7YnTCLVuw594s,3835
-xax/task/mixins/checkpointing.py,sha256=AMlobojybvJdDZcNCxm1DHSVC_2Qvnu_MbRcsc_8eoA,8508
+xax/task/mixins/artifacts.py,sha256=2ezmZGzPGe3nhsd9KRkeHWWXdbT9m7drzimIfw6v1XY,2892
+xax/task/mixins/checkpointing.py,sha256=sRkVxJbQfqDf1-lp1KFrAGYWHhTlV8_DORxGQ_69P1A,8954
 xax/task/mixins/compile.py,sha256=FRsxwLnZjjxpeWJ7Bx_d8XUY50oDoGidgpeRt4ejeQk,3377
 xax/task/mixins/cpu_stats.py,sha256=C_t71UTrv4LwQzhO5iubsfomj4jYa9bzpE4zBcHdoHM,9211
 xax/task/mixins/data_loader.py,sha256=WjMWk9uACfBMMClLMcLPkE0WNIvlCZnmqyyqLqJpjX0,6545
 xax/task/mixins/gpu_stats.py,sha256=IGPBro9xzSivwD43zM18lWcuei7IhA8LilxSPHqNl4I,8747
-xax/task/mixins/logger.py,sha256=CIQ4w4K3FcxN6A9xUfITdVkulSxPa4iaTe6cbs9ruaM,1958
+xax/task/mixins/logger.py,sha256=6oXsJJyNUx6YT3q58FVXMZBUpMgjVkGre6BXFN20cVI,2808
 xax/task/mixins/process.py,sha256=d1opVgvc6bOFXb7R58b07F4P5lbSZIzYaajtE0eBbpw,1477
 xax/task/mixins/runnable.py,sha256=IYIsLd2k09g-_y6o44EhJqT7E6BpsyEMmsyLSuzqjtc,1979
-xax/task/mixins/step_wrapper.py,sha256=DJw42mUGwgKx2tkeqatKR9_F4J8ug4wmxKMeJPmhcVQ,1560
-xax/task/mixins/train.py,sha256=dhGL_IuDaJy39BooYlO7JO-_EotKldtBhBplDGU_AnM,21745
+xax/task/mixins/step_wrapper.py,sha256=-Yu5Nft2CRw1JvZt6J_94SM1vqX8fk08IDK95Pmd2ew,1648
+xax/task/mixins/train.py,sha256=BEC7HSwBlGZDe7jCsedqEA8-K1Zx52-bTjsBONYIE5g,22225
 xax/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-xax/utils/experiments.py,sha256=qT3H0fyVH8DN417x7T0Xmz4SKoogW81-EHcZfyktFI8,28300
+xax/utils/debugging.py,sha256=9WlCrEqbq-SVXPEM4rhsLYERH97XNX7XSYLSI3sgKGk,1619
+xax/utils/experiments.py,sha256=_cwoBaiBxoQ_Tstm0rz7TEqfELqcktmPflb6AP1K0qA,28779
 xax/utils/jax.py,sha256=tC0NNelbrSTzwNGluiwLGKtoHhVpgdzrv-xherB3VtY,4752
-xax/utils/logging.py,sha256=ST1hp2C2xntVVJBUHwo3YxPK19fBLNvHU2WGO1xqcXA,6418
+xax/utils/jaxpr.py,sha256=S80nyEkv188RInzq3kCAdkQCU-bf6s0oPTrCE_LjkRs,2298
+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=Jwx6ErJfv1r2D23D4eKz1Hoo3mAJ0SEqC3EagZarWkw,1858
-xax/utils/tensorboard.py,sha256=oGq2E3Yr0z2xaACv2UOVt_CHEVc8fBxI8V1M99Fd34E,9742
+xax/utils/pytree.py,sha256=7GjQoPc_ZSZt3QS_9qXoBWl1jfMp1qZa7aViQoWJ0OQ,8864
+xax/utils/tensorboard.py,sha256=_S70dS69pduiD05viHAGgYGsaBry1QL2ej6ZwUIXPOE,16170
 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.0.7.dist-info/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
-xax-0.0.7.dist-info/METADATA,sha256=hE0KO4kYcN6Ed8iZ4649R5ENOUaQysBMW9vTh-94d4I,1171
-xax-0.0.7.dist-info/WHEEL,sha256=52BFRY2Up02UkjOa29eZOS2VxUrpPORXg1pkohGGUS8,91
-xax-0.0.7.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
-xax-0.0.7.dist-info/RECORD,,
+xax-0.1.1.dist-info/licenses/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
+xax-0.1.1.dist-info/METADATA,sha256=tJ4ilL3uBbykHBQTHbh-bN6m4hrHqivyyFeuI33ddX4,1877
+xax-0.1.1.dist-info/WHEEL,sha256=1tXe9gY0PYatrMPMDd6jXqjfpz_B-Wqm32CPfRC58XU,91
+xax-0.1.1.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
+xax-0.1.1.dist-info/RECORD,,

{xax-0.0.7.dist-info → xax-0.1.1.dist-info}/WHEEL

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