xax 0.1.4__py3-none-any.whl → 0.1.5__py3-none-any.whl

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.5"
 
 # This list shouldn't be modified by hand; instead, run the update script.
 __all__ = [

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/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/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.dist-info → xax-0.1.5.dist-info}/METADATA

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

{xax-0.1.4.dist-info → xax-0.1.5.dist-info}/RECORD

@@ -1,4 +1,4 @@
-xax/__init__.py,sha256=Ib-tn4sRqSIiFjoQceuUzdwRc4i-VW9P5bBLq4wDbYc,13361
+xax/__init__.py,sha256=Ig2JALbHQOtFHr57jpmPm16L0aYkNTNZ-Upz9jxIyrc,13361
 xax/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/requirements-dev.txt,sha256=qkscNkFzWd1S5fump-AKH53rR65v2x5FmboFdy_kKvs,128
 xax/requirements.txt,sha256=9LAEZ5c5gqRSARRVA6xJsVTa4MebPZuC4yOkkwkZJFw,297
@@ -8,7 +8,7 @@ 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/export.py,sha256=bu2m-4FDnadEhXDb9zM6SgOZvsf5p4xiee1sFZyNF7c,5510
 xax/nn/functions.py,sha256=CI_OmspaQwN9nl4hwefIU3_I7m6gBZwJ9aGK1JGUgr0,2713
 xax/nn/geom.py,sha256=eK7I8fUHBc3FT7zpm5Yf__bXFQ4LtX6sa17-DxojLTo,3202
 xax/nn/norm.py,sha256=cDmYf5CtyzmuCiWdSP5nr8nZKQOmaZueDQXMPnThg6c,548
@@ -39,7 +39,7 @@ xax/task/mixins/logger.py,sha256=6oXsJJyNUx6YT3q58FVXMZBUpMgjVkGre6BXFN20cVI,280
 xax/task/mixins/process.py,sha256=d1opVgvc6bOFXb7R58b07F4P5lbSZIzYaajtE0eBbpw,1477
 xax/task/mixins/runnable.py,sha256=IYIsLd2k09g-_y6o44EhJqT7E6BpsyEMmsyLSuzqjtc,1979
 xax/task/mixins/step_wrapper.py,sha256=-Yu5Nft2CRw1JvZt6J_94SM1vqX8fk08IDK95Pmd2ew,1648
-xax/task/mixins/train.py,sha256=8AaBXaopnrxtSZXldyFCE3QX1k5r3IsZMr6O0ICnNnU,22332
+xax/task/mixins/train.py,sha256=vsH_QpyrThlh9AzWnyvDJv58Y8U_516oi8gmMq_0iMg,22333
 xax/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/utils/debugging.py,sha256=9WlCrEqbq-SVXPEM4rhsLYERH97XNX7XSYLSI3sgKGk,1619
 xax/utils/experiments.py,sha256=5CUja1H_cx4dnVqTGQekOpIhqISwHtAgLxZ34GV7cwM,29229
@@ -49,12 +49,12 @@ 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=7GjQoPc_ZSZt3QS_9qXoBWl1jfMp1qZa7aViQoWJ0OQ,8864
-xax/utils/tensorboard.py,sha256=_S70dS69pduiD05viHAGgYGsaBry1QL2ej6ZwUIXPOE,16170
+xax/utils/tensorboard.py,sha256=21czW8WC2SAmwEhz6RLJc_q5HFvNKM4iR1ZycSO5qPE,17058
 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.1.4.dist-info/licenses/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
-xax-0.1.4.dist-info/METADATA,sha256=VMyCF4gPJh5ww_VzF0HgfInf_bfD40ahwxpHkxEMGoI,1877
-xax-0.1.4.dist-info/WHEEL,sha256=DK49LOLCYiurdXXOXwGJm6U4DkHkg4lcxjhqwRa0CP4,91
-xax-0.1.4.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
-xax-0.1.4.dist-info/RECORD,,
+xax-0.1.5.dist-info/licenses/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
+xax-0.1.5.dist-info/METADATA,sha256=dkvPH-GXRErGzzR0VGsiRisNYBzDVf97nEGG0F2-HjY,1877
+xax-0.1.5.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
+xax-0.1.5.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
+xax-0.1.5.dist-info/RECORD,,

{xax-0.1.4.dist-info → xax-0.1.5.dist-info}/WHEEL

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