xax 0.2.7__tar.gz → 0.2.9__tar.gz

{xax-0.2.7/xax.egg-info → xax-0.2.9}/PKG-INFO

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

{xax-0.2.7 → xax-0.2.9}/pyproject.toml

@@ -57,6 +57,7 @@ target-version = "py311"
 
 select = [
     "ANN",
+    "B",
     "D",
     "E",
     "F",

{xax-0.2.7 → xax-0.2.9}/xax/__init__.py

@@ -12,7 +12,7 @@ and running the update script:
     python -m scripts.update_api --inplace
 """
 
-__version__ = "0.2.7"
+__version__ = "0.2.9"
 
 # This list shouldn't be modified by hand; instead, run the update script.
 __all__ = [
@@ -43,9 +43,12 @@ __all__ = [
     "cubic_bezier_interpolation",
     "euler_to_quat",
     "get_projected_gravity_vector_from_quat",
+    "normalize",
     "quat_to_euler",
     "quat_to_rotmat",
     "rotate_vector_by_quat",
+    "rotation6d_to_rotation_matrix",
+    "rotation_matrix_to_rotation6d",
     "cross_entropy",
     "cast_norm_type",
     "get_norm",
@@ -57,8 +60,18 @@ __all__ = [
     "BaseLauncher",
     "CliLauncher",
     "SingleProcessLauncher",
+    "LogDistribution",
+    "LogError",
+    "LogErrorSummary",
+    "LogGraph",
+    "LogHistogram",
     "LogImage",
     "LogLine",
+    "LogMesh",
+    "LogPing",
+    "LogScalar",
+    "LogStatus",
+    "LogVideo",
     "Logger",
     "LoggerImpl",
     "CallbackLogger",
@@ -72,7 +85,6 @@ __all__ = [
     "GPUStatsOptions",
     "StepContext",
     "ValidStepTimer",
-    "get_param_count",
     "Script",
     "ScriptConfig",
     "Config",
@@ -117,6 +129,7 @@ __all__ = [
     "compute_nan_ratio",
     "flatten_array",
     "flatten_pytree",
+    "get_pytree_param_count",
     "pytree_has_nans",
     "reshuffle_pytree",
     "reshuffle_pytree_along_dims",
@@ -209,9 +222,12 @@ NAME_MAP: dict[str, str] = {
     "cubic_bezier_interpolation": "nn.geom",
     "euler_to_quat": "nn.geom",
     "get_projected_gravity_vector_from_quat": "nn.geom",
+    "normalize": "nn.geom",
     "quat_to_euler": "nn.geom",
     "quat_to_rotmat": "nn.geom",
     "rotate_vector_by_quat": "nn.geom",
+    "rotation6d_to_rotation_matrix": "nn.geom",
+    "rotation_matrix_to_rotation6d": "nn.geom",
     "cross_entropy": "nn.losses",
     "cast_norm_type": "nn.norm",
     "get_norm": "nn.norm",
@@ -223,8 +239,18 @@ NAME_MAP: dict[str, str] = {
     "BaseLauncher": "task.launchers.base",
     "CliLauncher": "task.launchers.cli",
     "SingleProcessLauncher": "task.launchers.single_process",
+    "LogDistribution": "task.logger",
+    "LogError": "task.logger",
+    "LogErrorSummary": "task.logger",
+    "LogGraph": "task.logger",
+    "LogHistogram": "task.logger",
     "LogImage": "task.logger",
     "LogLine": "task.logger",
+    "LogMesh": "task.logger",
+    "LogPing": "task.logger",
+    "LogScalar": "task.logger",
+    "LogStatus": "task.logger",
+    "LogVideo": "task.logger",
     "Logger": "task.logger",
     "LoggerImpl": "task.logger",
     "CallbackLogger": "task.loggers.callback",
@@ -238,7 +264,6 @@ NAME_MAP: dict[str, str] = {
     "GPUStatsOptions": "task.mixins.gpu_stats",
     "StepContext": "task.mixins.step_wrapper",
     "ValidStepTimer": "task.mixins.train",
-    "get_param_count": "task.mixins.train",
     "Script": "task.script",
     "ScriptConfig": "task.script",
     "Config": "task.task",
@@ -283,6 +308,7 @@ NAME_MAP: dict[str, str] = {
     "compute_nan_ratio": "utils.pytree",
     "flatten_array": "utils.pytree",
     "flatten_pytree": "utils.pytree",
+    "get_pytree_param_count": "utils.pytree",
     "pytree_has_nans": "utils.pytree",
     "reshuffle_pytree": "utils.pytree",
     "reshuffle_pytree_along_dims": "utils.pytree",
@@ -376,9 +402,12 @@ if IMPORT_ALL or TYPE_CHECKING:
         cubic_bezier_interpolation,
         euler_to_quat,
         get_projected_gravity_vector_from_quat,
+        normalize,
         quat_to_euler,
         quat_to_rotmat,
         rotate_vector_by_quat,
+        rotation6d_to_rotation_matrix,
+        rotation_matrix_to_rotation6d,
     )
     from xax.nn.losses import cross_entropy
     from xax.nn.norm import NormType, cast_norm_type, get_norm
@@ -388,7 +417,22 @@ if IMPORT_ALL or TYPE_CHECKING:
     from xax.task.launchers.base import BaseLauncher
     from xax.task.launchers.cli import CliLauncher
     from xax.task.launchers.single_process import SingleProcessLauncher
-    from xax.task.logger import Logger, LoggerImpl, LogImage, LogLine
+    from xax.task.logger import (
+        LogDistribution,
+        LogError,
+        LogErrorSummary,
+        Logger,
+        LoggerImpl,
+        LogGraph,
+        LogHistogram,
+        LogImage,
+        LogLine,
+        LogMesh,
+        LogPing,
+        LogScalar,
+        LogStatus,
+        LogVideo,
+    )
     from xax.task.loggers.callback import CallbackLogger
     from xax.task.loggers.json import JsonLogger
     from xax.task.loggers.state import StateLogger
@@ -399,7 +443,7 @@ if IMPORT_ALL or TYPE_CHECKING:
     from xax.task.mixins.data_loader import DataloaderConfig
     from xax.task.mixins.gpu_stats import GPUStatsOptions
     from xax.task.mixins.step_wrapper import StepContext
-    from xax.task.mixins.train import Batch, Output, ValidStepTimer, get_param_count
+    from xax.task.mixins.train import Batch, Output, ValidStepTimer
     from xax.task.script import Script, ScriptConfig
     from xax.task.task import Config, Task
     from xax.utils.data.collate import CollateMode, collate, collate_non_null
@@ -444,6 +488,7 @@ if IMPORT_ALL or TYPE_CHECKING:
         compute_nan_ratio,
         flatten_array,
         flatten_pytree,
+        get_pytree_param_count,
         pytree_has_nans,
         reshuffle_pytree,
         reshuffle_pytree_along_dims,

{xax-0.2.7 → xax-0.2.9}/xax/core/conf.py

@@ -26,7 +26,7 @@ def field(value: FieldType, **kwargs: str) -> FieldType:
     metadata: dict[str, Any] = {}
     metadata.update(kwargs)
 
-    if hasattr(value, "__call__"):
+    if hasattr(value, "__call__"):  # noqa: B004
         return field_base(default_factory=value, metadata=metadata)
     if value.__class__.__hash__ is None:
         return field_base(default_factory=lambda: value, metadata=metadata)

{xax-0.2.7 → xax-0.2.9}/xax/nn/equinox.py

@@ -68,8 +68,8 @@ def _infer_activation(activation: ActivationFunction) -> Callable:
         return lambda x: x
     try:
         return getattr(jax.nn, activation)
-    except AttributeError:
-        raise ValueError(f"Activation function `{activation}` not found in `jax.nn`")
+    except AttributeError as err:
+        raise ValueError(f"Activation function `{activation}` not found in `jax.nn`") from err
 
 
 def make_eqx_mlp(hyperparams: MLPHyperParams, *, key: PRNGKeyArray) -> eqx.nn.MLP:
@@ -100,7 +100,7 @@ def make_eqx_mlp(hyperparams: MLPHyperParams, *, key: PRNGKeyArray) -> eqx.nn.ML
 def export_eqx_mlp(
     model: eqx.nn.MLP,
     output_path: str | Path,
-    dtype: jax.numpy.dtype = eqx._misc.default_floating_dtype(),
+    dtype: jax.numpy.dtype | None = None,
 ) -> None:
     """Serialize an Equinox MLP to a .eqx file.
 
@@ -109,6 +109,9 @@ def export_eqx_mlp(
         output_path: The path to save the exported model.
         dtype: The dtype of the model.
     """
+    if dtype is None:
+        dtype = eqx._misc.default_floating_dtype()
+
     activation = model.activation.__name__
     final_activation = model.final_activation.__name__
 

{xax-0.2.7 → xax-0.2.9}/xax/nn/functions.py

@@ -58,13 +58,16 @@ def recursive_chunk(item: Any, num_chunks: int, dim: int = 0) -> Iterable[Any]:
         yield from np.array_split(item, num_chunks, axis=dim)
     elif is_dataclass(item):
         yield from (
-            item.__class__(**{k: i for k, i in zip(item.__dict__, ii)})
-            for ii in zip(*(recursive_chunk(v, num_chunks, dim) for v in item.__dict__.values()))
+            item.__class__(**{k: i for k, i in zip(item.__dict__, ii, strict=True)})
+            for ii in zip(*(recursive_chunk(v, num_chunks, dim) for v in item.__dict__.values()), strict=False)
         )
     elif isinstance(item, Mapping):
-        yield from (dict(zip(item, ii)) for ii in zip(*(recursive_chunk(i, num_chunks, dim) for i in item.values())))
+        yield from (
+            dict(zip(item, ii, strict=False))
+            for ii in zip(*(recursive_chunk(i, num_chunks, dim) for i in item.values()), strict=False)
+        )
     elif isinstance(item, Sequence):
-        yield from (list(ii) for ii in zip(*(recursive_chunk(i, num_chunks, dim) for i in item)))
+        yield from (list(ii) for ii in zip(*(recursive_chunk(i, num_chunks, dim) for i in item), strict=False))
     else:
         yield from (item for _ in range(num_chunks))
 

{xax-0.2.7 → xax-0.2.9}/xax/nn/geom.py

@@ -1,5 +1,6 @@
 """Defines geometry functions."""
 
+import chex
 from jax import numpy as jnp
 from jaxtyping import Array
 
@@ -211,3 +212,51 @@ def quat_to_rotmat(quat: Array, eps: float = 1e-6) -> Array:
         ],
         axis=-2,
     )
+
+
+def normalize(v: jnp.ndarray, axis: int = -1, eps: float = 1e-8) -> jnp.ndarray:
+    norm = jnp.linalg.norm(v, axis=axis, keepdims=True)
+    return v / jnp.clip(norm, a_min=eps)
+
+
+def rotation6d_to_rotation_matrix(r6d: jnp.ndarray) -> jnp.ndarray:
+    """Convert 6D rotation representation to rotation matrix.
+
+    From https://arxiv.org/pdf/1812.07035, Appendix B
+
+    Args:
+        r6d: The 6D rotation representation, shape (*, 6).
+
+    Returns:
+        The rotation matrix, shape (*, 3, 3).
+    """
+    chex.assert_shape(r6d, (..., 6))
+    shape = r6d.shape
+    flat = r6d.reshape(-1, 6)
+    a_1 = flat[:, 0:3]
+    a_2 = flat[:, 3:6]
+
+    b_1 = normalize(a_1, axis=-1)
+
+    # Reordered Gram-Schmidt orthonormalization.
+    b_3 = normalize(jnp.cross(b_1, a_2), axis=-1)
+    b_2 = jnp.cross(b_3, b_1)
+
+    rotation_matrix = jnp.stack([b_1, b_2, b_3], axis=-1)
+    return rotation_matrix.reshape(shape[:-1] + (3, 3))
+
+
+def rotation_matrix_to_rotation6d(rotation_matrix: jnp.ndarray) -> jnp.ndarray:
+    """Convert rotation matrix to 6D rotation representation.
+
+    Args:
+        rotation_matrix: The rotation matrix, shape (*, 3, 3).
+
+    Returns:
+        The 6D rotation representation, shape (*, 6).
+    """
+    chex.assert_shape(rotation_matrix, (..., 3, 3))
+    shape = rotation_matrix.shape
+    # Simply concatenate a1 and a2 from SO(3)
+    r6d = jnp.concatenate([rotation_matrix[..., 0], rotation_matrix[..., 1]], axis=-1)
+    return r6d.reshape(shape[:-2] + (6,))

{xax-0.2.7 → xax-0.2.9}/xax/task/base.py

@@ -184,8 +184,8 @@ class BaseTask(Generic[Config]):
 
             # Attempts to load any paths as configs.
             is_path = [Path(arg).is_file() or (task_path / arg).is_file() for arg in args]
-            paths = [arg for arg, is_path in zip(args, is_path) if is_path]
-            non_paths = [arg for arg, is_path in zip(args, is_path) if not is_path]
+            paths = [arg for arg, is_path in zip(args, is_path, strict=True) if is_path]
+            non_paths = [arg for arg, is_path in zip(args, is_path, strict=True) if not is_path]
             if paths:
                 cfg = OmegaConf.merge(cfg, *(get_config(path, task_path) for path in paths))
             cfg = OmegaConf.merge(cfg, OmegaConf.from_cli(non_paths))

{xax-0.2.7 → xax-0.2.9}/xax/task/logger.py

@@ -462,11 +462,11 @@ class LoggerImpl(ABC):
 
         self.tickers = {phase: IntervalTicker(log_interval_seconds) for phase in get_args(Phase)}
 
-    def start(self) -> None:
-        pass
+    @abstractmethod
+    def start(self) -> None: ...
 
-    def stop(self) -> None:
-        pass
+    @abstractmethod
+    def stop(self) -> None: ...
 
     @abstractmethod
     def write(self, line: LogLine) -> None:
@@ -476,6 +476,7 @@ class LoggerImpl(ABC):
             line: The line to write.
         """
 
+    @abstractmethod
     def write_error_summary(self, error_summary: LogErrorSummary) -> None:
         """Handles writing an error summary.
 
@@ -483,6 +484,7 @@ class LoggerImpl(ABC):
             error_summary: The error summary to write.
         """
 
+    @abstractmethod
     def write_error(self, error: LogError) -> None:
         """Handles writing an error line.
 
@@ -490,6 +492,7 @@ class LoggerImpl(ABC):
             error: The error information to write.
         """
 
+    @abstractmethod
     def write_status(self, status: LogStatus) -> None:
         """Handles writing a status line.
 
@@ -497,6 +500,7 @@ class LoggerImpl(ABC):
             status: The status to write.
         """
 
+    @abstractmethod
     def write_ping(self, ping: LogPing) -> None:
         """Handles writing a ping line.
 
@@ -504,6 +508,7 @@ class LoggerImpl(ABC):
             ping: The ping to write.
         """
 
+    @abstractmethod
     def log_file(self, name: str, contents: str) -> None:
         """Logs a large text file.
 
@@ -621,7 +626,7 @@ class Logger:
             return
         line = self.pack(state)
         self.clear()
-        for lg in (lg for lg, should_log in zip(self.loggers, should_log) if should_log):
+        for lg in (lg for lg, should_log in zip(self.loggers, should_log, strict=False) if should_log):
             lg.write(line)
 
     def write_error_summary(self, error_summary: str) -> None:
@@ -1045,7 +1050,7 @@ class Logger:
                         line_spacing=line_spacing,
                         centered=centered,
                     )
-                    for img, label in zip(images, labels)
+                    for img, label in zip(images, labels, strict=True)
                 ]
                 tiled = tile_images([img.image for img in labeled], sep)
 

{xax-0.2.7 → xax-0.2.9}/xax/task/loggers/callback.py

@@ -25,6 +25,12 @@ class CallbackLogger(LoggerImpl):
         self.ping_callback = ping_callback
         self.file_callback = file_callback
 
+    def start(self) -> None:
+        pass
+
+    def stop(self) -> None:
+        pass
+
     def write(self, line: LogLine) -> None:
         self.callback(line)
 

{xax-0.2.7 → xax-0.2.9}/xax/task/loggers/json.py

@@ -8,6 +8,7 @@ from jaxtyping import Array
 
 from xax.task.logger import (
     LogError,
+    LogErrorSummary,
     LoggerImpl,
     LogLine,
     LogPing,
@@ -57,6 +58,12 @@ class JsonLogger(LoggerImpl):
         self.line_sep = line_sep
         self.remove_unicode_from_namespaces = remove_unicode_from_namespaces
 
+    def start(self) -> None:
+        pass
+
+    def stop(self) -> None:
+        pass
+
     @property
     def fp(self) -> TextIO:
         return self.log_stream
@@ -87,6 +94,12 @@ class JsonLogger(LoggerImpl):
         if self.flush_immediately:
             self.fp.flush()
 
+    def write_error_summary(self, error_summary: LogErrorSummary) -> None:
+        pass
+
+    def log_file(self, name: str, contents: str) -> None:
+        pass
+
     def write_error(self, error: LogError) -> None:
         self.err_fp.write(error.message)
         if error.location is not None:

{xax-0.2.7 → xax-0.2.9}/xax/task/loggers/state.py

@@ -3,7 +3,14 @@
 from pathlib import Path
 from typing import Literal
 
-from xax.task.logger import LoggerImpl, LogLine
+from xax.task.logger import (
+    LogError,
+    LogErrorSummary,
+    LoggerImpl,
+    LogLine,
+    LogPing,
+    LogStatus,
+)
 
 
 class StateLogger(LoggerImpl):
@@ -30,3 +37,21 @@ class StateLogger(LoggerImpl):
 
     def write(self, line: LogLine) -> None:
         pass
+
+    def start(self) -> None:
+        pass
+
+    def stop(self) -> None:
+        pass
+
+    def write_error_summary(self, error_summary: LogErrorSummary) -> None:
+        pass
+
+    def write_error(self, error: LogError) -> None:
+        pass
+
+    def write_status(self, status: LogStatus) -> None:
+        pass
+
+    def write_ping(self, ping: LogPing) -> None:
+        pass

{xax-0.2.7 → xax-0.2.9}/xax/task/loggers/stdout.py

@@ -79,11 +79,13 @@ class StdoutLogger(LoggerImpl):
         self.error_summary: tuple[str, datetime.datetime] | None = None
 
     def start(self) -> None:
-        return super().start()
+        pass
 
     def stop(self) -> None:
         self.write_queues()
-        return super().stop()
+
+    def log_file(self, name: str, contents: str) -> None:
+        pass
 
     def write_separator(self) -> None:
         self.write_fp.write("\033[2J\033[H")

{xax-0.2.7 → xax-0.2.9}/xax/task/loggers/tensorboard.py

@@ -12,7 +12,7 @@ from typing import TypeVar
 
 from xax.core.state import Phase
 from xax.nn.parallel import is_master
-from xax.task.logger import LoggerImpl, LogLine
+from xax.task.logger import LogError, LogErrorSummary, LoggerImpl, LogLine, LogPing, LogStatus
 from xax.utils.jax import as_float
 from xax.utils.logging import LOG_STATUS, port_is_busy
 from xax.utils.tensorboard import TensorboardWriter, TensorboardWriters
@@ -236,3 +236,21 @@ class TensorboardLogger(LoggerImpl):
         for name, contents in self.files.items():
             writer.add_text(name, contents)
         self.files.clear()
+
+    def start(self) -> None:
+        pass
+
+    def stop(self) -> None:
+        pass
+
+    def write_error(self, error: LogError) -> None:
+        pass
+
+    def write_error_summary(self, error_summary: LogErrorSummary) -> None:
+        pass
+
+    def write_ping(self, ping: LogPing) -> None:
+        pass
+
+    def write_status(self, status: LogStatus) -> None:
+        pass

{xax-0.2.7 → xax-0.2.9}/xax/task/mixins/artifacts.py

@@ -31,11 +31,13 @@ Config = TypeVar("Config", bound=ArtifactsConfig)
 
 class ArtifactsMixin(BaseTask[Config]):
     _exp_dir: Path | None
+    _stage_dir: Path | None
 
     def __init__(self, config: Config) -> None:
         super().__init__(config)
 
         self._exp_dir = None
+        self._stage_dir = None
 
     @functools.cached_property
     def run_dir(self) -> Path:
@@ -75,15 +77,16 @@ class ArtifactsMixin(BaseTask[Config]):
         logger.log(LOG_STATUS, self._exp_dir)
         return self._exp_dir
 
-    @functools.lru_cache(maxsize=None)
     def stage_environment(self) -> Path | None:
-        stage_dir = (self.exp_dir / "code").resolve()
-        try:
-            stage_environment(self, stage_dir)
-        except Exception:
-            logger.exception("Failed to stage environment!")
-            return None
-        return stage_dir
+        if self._stage_dir is None:
+            stage_dir = (self.exp_dir / "code").resolve()
+            try:
+                stage_environment(self, stage_dir)
+            except Exception:
+                logger.exception("Failed to stage environment!")
+                return None
+            self._stage_dir = stage_dir
+        return self._stage_dir
 
     def on_training_end(self, state: State) -> State:
         state = super().on_training_end(state)

{xax-0.2.7 → xax-0.2.9}/xax/task/mixins/train.py

@@ -57,6 +57,7 @@ from xax.utils.experiments import (
 )
 from xax.utils.jax import jit as xax_jit
 from xax.utils.logging import LOG_PING, LOG_STATUS
+from xax.utils.pytree import get_pytree_param_count
 from xax.utils.text import highlight_exception_message, show_info
 from xax.utils.types.frozen_dict import FrozenDict
 
@@ -96,12 +97,6 @@ def batches_per_step_schedule(schedule: list[int] | None) -> list[int] | None:
     return list(itertools.accumulate([0] + schedule))
 
 
-def get_param_count(pytree: PyTree) -> int:
-    """Calculates the total number of parameters in a PyTree."""
-    leaves, _ = jax.tree.flatten(pytree)
-    return sum(x.size for x in leaves if isinstance(x, jnp.ndarray))
-
-
 class ValidStepTimer:
     def __init__(
         self,
@@ -690,7 +685,7 @@ class TrainMixin(
         self.logger.log_file("info.json", get_info_json())
 
     def log_model_size(self, model: PyTree) -> None:
-        logger.info("Model size: %s", f"{get_param_count(model):,}")
+        logger.info("Model size: %s", f"{get_pytree_param_count(model):,}")
 
     def model_partition_fn(self, item: Any) -> bool:  # noqa: ANN401
         return eqx.is_inexact_array(item)

{xax-0.2.7 → xax-0.2.9}/xax/utils/experiments.py

@@ -749,7 +749,8 @@ class BaseFileDownloader(ABC):
                         f"We detected some HTML elements in the downloaded file. "
                         f"This most likely means that the download triggered an unhandled API response by GDrive. "
                         f"Please report this to torchvision at https://github.com/pytorch/vision/issues including "
-                        f"the response:\n\n{text}"
+                        f"the response:\n\n{text}",
+                        stacklevel=2,
                     )
 
     @classmethod

{xax-0.2.7 → xax-0.2.9}/xax/utils/pytree.py

@@ -1,6 +1,7 @@
 """Utils for accessing, modifying, and otherwise manipulating pytrees."""
 
 import chex
+import equinox as eqx
 import jax
 import jax.numpy as jnp
 from jax import Array
@@ -124,7 +125,7 @@ def reshuffle_pytree(data: PyTree, batch_shape: tuple[int, ...], rng: PRNGKeyArr
 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)]
+    perms = [jax.random.permutation(rng_i, dim) for rng_i, dim in zip(rngs, batch_shape, strict=True)]
     # n-dimensional index grid from permutations
     idx_grids = jnp.meshgrid(*perms, indexing="ij")
 
@@ -236,3 +237,9 @@ def reshuffle_pytree_along_dims(
         return x
 
     return jax.tree.map_with_path(restore_transpose, reshuffled_transposed)
+
+
+def get_pytree_param_count(pytree: PyTree) -> int:
+    """Calculates the total number of parameters in a PyTree."""
+    leaves, _ = jax.tree.flatten(pytree)
+    return sum(x.size for x in leaves if isinstance(x, jnp.ndarray) and eqx.is_inexact_array(x))

{xax-0.2.7 → xax-0.2.9}/xax/utils/text.py

@@ -192,7 +192,7 @@ def render_text_blocks(
         if any(len(row) != len(blocks[0]) for row in blocks):
             raise ValueError("All rows must have the same number of blocks in order to align them")
         widths = [[max(len(line) for line in i.lines) if i.width is None else i.width for i in r] for r in blocks]
-        row_widths = [max(i) for i in zip(*widths)]
+        row_widths = [max(i) for i in zip(*widths, strict=True)]
         for row in blocks:
             for i, block in enumerate(row):
                 block.width = row_widths[i]
@@ -263,7 +263,7 @@ def render_text_blocks(
                             if i >= len(block.lines)
                             else colored(pad(block.lines[i], width, block.center), block.color, bold=block.bold)
                         )
-                        for block, width in zip(row, get_widths(row))
+                        for block, width in zip(row, get_widths(row), strict=True)
                     ]
                 )
                 + " │"

{xax-0.2.7 → xax-0.2.9}/xax/utils/types/frozen_dict.py

@@ -133,7 +133,7 @@ class FrozenDict(Mapping[K, V]):
 
     @classmethod
     def tree_unflatten(cls, keys: tuple[K, ...], values: tuple[Any, ...]) -> "FrozenDict[K, V]":
-        return cls({k: v for k, v in zip(keys, values)}, __unsafe_skip_copy__=True)
+        return cls({k: v for k, v in zip(keys, values, strict=True)}, __unsafe_skip_copy__=True)
 
 
 def unfreeze(x: FrozenDict[K, V] | dict[str, Any]) -> dict[Any, Any]:  # noqa: ANN401

{xax-0.2.7 → xax-0.2.9/xax.egg-info}/PKG-INFO

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