xax 0.1.4__tar.gz → 0.1.6__tar.gz

{xax-0.1.4/xax.egg-info → xax-0.1.6}/PKG-INFO

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

{xax-0.1.4 → xax-0.1.6}/xax/__init__.py

@@ -12,7 +12,7 @@ and running the update script:
     python -m scripts.update_api --inplace
 """
 
-__version__ = "0.1.4"
+__version__ = "0.1.6"
 
 # This list shouldn't be modified by hand; instead, run the update script.
 __all__ = [
@@ -97,6 +97,8 @@ __all__ = [
     "save_config",
     "stage_environment",
     "to_markdown_table",
+    "HashableArray",
+    "hashable_array",
     "jit",
     "save_jaxpr_dot",
     "ColoredFormatter",
@@ -251,6 +253,8 @@ NAME_MAP: dict[str, str] = {
     "save_config": "utils.experiments",
     "stage_environment": "utils.experiments",
     "to_markdown_table": "utils.experiments",
+    "HashableArray": "utils.jax",
+    "hashable_array": "utils.jax",
     "jit": "utils.jax",
     "save_jaxpr_dot": "utils.jaxpr",
     "ColoredFormatter": "utils.logging",
@@ -404,7 +408,7 @@ if IMPORT_ALL or TYPE_CHECKING:
         stage_environment,
         to_markdown_table,
     )
-    from xax.utils.jax import jit
+    from xax.utils.jax import HashableArray, hashable_array, jit
     from xax.utils.jaxpr import save_jaxpr_dot
     from xax.utils.logging import (
         LOG_ERROR_SUMMARY,

{xax-0.1.4 → xax-0.1.6}/xax/nn/export.py

@@ -9,7 +9,14 @@ import jax
 import tensorflow as tf
 from jax.experimental import jax2tf
 from jaxtyping import Array, PyTree
-from orbax.export import ExportManager, JaxModule, ServingConfig
+
+try:
+    from orbax.export import ExportManager, JaxModule, ServingConfig
+except ImportError as e:
+    raise ImportError(
+        "Please install the package with `orbax` as a dependency, using "
+        "'xax[export]` to install the required dependencies."
+    ) from e
 
 logger = logging.getLogger(__name__)
 

{xax-0.1.4 → xax-0.1.6}/xax/task/mixins/train.py

@@ -122,7 +122,7 @@ class ValidStepTimer:
 
         # Step-based validation.
         valid_every_n_steps = self.valid_every_n_steps
-        if valid_every_n_steps is not None and state.num_steps > valid_every_n_steps + self.last_valid_step:
+        if valid_every_n_steps is not None and state.num_steps >= valid_every_n_steps + self.last_valid_step:
             self.last_valid_step = state.num_steps
             return True
 

{xax-0.1.4 → xax-0.1.6}/xax/utils/jax.py

@@ -138,3 +138,25 @@ def jit(
         return wrapped
 
     return decorator
+
+
+class HashableArray:
+    def __init__(self, array: np.ndarray | jnp.ndarray) -> None:
+        if not isinstance(array, (np.ndarray, jnp.ndarray)):
+            raise ValueError(f"Expected np.ndarray or jnp.ndarray, got {type(array)}")
+        self.array = array
+        self._hash: int | None = None
+
+    def __hash__(self) -> int:
+        if self._hash is None:
+            self._hash = hash(self.array.tobytes())
+        return self._hash
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, HashableArray):
+            return False
+        return bool(jnp.array_equal(self.array, other.array))
+
+
+def hashable_array(array: np.ndarray | jnp.ndarray) -> HashableArray:
+    return HashableArray(array)

{xax-0.1.4 → xax-0.1.6}/xax/utils/tensorboard.py

@@ -258,12 +258,40 @@ class TensorboardWriter:
         fps: int = 30,
     ) -> None:
         assert value.ndim == 4, "Video must be 4D array (T, H, W, C)"
-        images = [PIL.Image.fromarray(frame) for frame in value]
+
+        images = [PIL.Image.fromarray(frame).convert("RGB") for frame in value]
+        width, height = images[0].size
+        big_image = PIL.Image.new("RGB", (width, height * len(images)))
+        for i, im in enumerate(images):
+            big_image.paste(im, (0, i * height))
+
+        quantized_big = big_image.quantize(method=PIL.Image.Quantize.MAXCOVERAGE, dither=PIL.Image.Dither.NONE)
+        palette = quantized_big.getpalette()
+
+        processed = []
+        for im in images:
+            q = im.quantize(
+                method=PIL.Image.Quantize.MAXCOVERAGE,
+                palette=quantized_big,
+                dither=PIL.Image.Dither.NONE,
+            )
+            processed.append(q)
+
+        if palette is not None:
+            palette[0:3] = [255, 255, 255]
+            for im in processed:
+                im.putpalette(palette)
 
         # Create temporary file for GIF
         temp_file = tempfile.NamedTemporaryFile(suffix=".gif", delete=False)
         try:
-            images[0].save(temp_file.name, save_all=True, append_images=images[1:], duration=int(1000 / fps), loop=0)
+            processed[0].save(
+                temp_file.name,
+                save_all=True,
+                append_images=processed[1:],
+                duration=int(1000 / fps),
+                loop=0,
+            )
             with open(temp_file.name, "rb") as f:
                 video_string = f.read()
 

{xax-0.1.4 → xax-0.1.6/xax.egg-info}/PKG-INFO

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