xax 0.1.13__py3-none-any.whl → 0.1.15__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.13"
+__version__ = "0.1.15"
 
 # This list shouldn't be modified by hand; instead, run the update script.
 __all__ = [
@@ -40,6 +40,7 @@ __all__ = [
     "load_eqx_mlp",
     "make_eqx_mlp",
     "save_eqx",
+    "cubic_bezier_interpolation",
     "euler_to_quat",
     "get_projected_gravity_vector_from_quat",
     "quat_to_euler",
@@ -77,6 +78,7 @@ __all__ = [
     "collate_non_null",
     "breakpoint_if_nan",
     "get_named_leaves",
+    "log_if_nan",
     "BaseFileDownloader",
     "ContextTimer",
     "CumulativeTimer",
@@ -200,6 +202,7 @@ NAME_MAP: dict[str, str] = {
     "load_eqx_mlp": "nn.equinox",
     "make_eqx_mlp": "nn.equinox",
     "save_eqx": "nn.equinox",
+    "cubic_bezier_interpolation": "nn.geom",
     "euler_to_quat": "nn.geom",
     "get_projected_gravity_vector_from_quat": "nn.geom",
     "quat_to_euler": "nn.geom",
@@ -237,6 +240,7 @@ NAME_MAP: dict[str, str] = {
     "collate_non_null": "utils.data.collate",
     "breakpoint_if_nan": "utils.debugging",
     "get_named_leaves": "utils.debugging",
+    "log_if_nan": "utils.debugging",
     "BaseFileDownloader": "utils.experiments",
     "ContextTimer": "utils.experiments",
     "CumulativeTimer": "utils.experiments",
@@ -361,6 +365,7 @@ if IMPORT_ALL or TYPE_CHECKING:
         save_eqx,
     )
     from xax.nn.geom import (
+        cubic_bezier_interpolation,
         euler_to_quat,
         get_projected_gravity_vector_from_quat,
         quat_to_euler,
@@ -388,7 +393,7 @@ if IMPORT_ALL or TYPE_CHECKING:
     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
-    from xax.utils.debugging import breakpoint_if_nan, get_named_leaves
+    from xax.utils.debugging import breakpoint_if_nan, get_named_leaves, log_if_nan
     from xax.utils.experiments import (
         BaseFileDownloader,
         ContextTimer,

xax/nn/geom.py

@@ -1,10 +1,10 @@
 """Defines geometry functions."""
 
-import jax
 from jax import numpy as jnp
+from jaxtyping import Array
 
 
-def quat_to_euler(quat_4: jax.Array, eps: float = 1e-6) -> jax.Array:
+def quat_to_euler(quat_4: Array, eps: float = 1e-6) -> Array:
     """Normalizes and converts a quaternion (w, x, y, z) to roll, pitch, yaw.
 
     Args:
@@ -40,7 +40,7 @@ def quat_to_euler(quat_4: jax.Array, eps: float = 1e-6) -> jax.Array:
     return jnp.concatenate([roll, pitch, yaw], axis=-1)
 
 
-def euler_to_quat(euler_3: jax.Array) -> jax.Array:
+def euler_to_quat(euler_3: Array) -> Array:
     """Converts roll, pitch, yaw angles to a quaternion (w, x, y, z).
 
     Args:
@@ -75,7 +75,7 @@ def euler_to_quat(euler_3: jax.Array) -> jax.Array:
     return quat
 
 
-def get_projected_gravity_vector_from_quat(quat: jax.Array, eps: float = 1e-6) -> jax.Array:
+def get_projected_gravity_vector_from_quat(quat: Array, eps: float = 1e-6) -> Array:
     """Calculates the gravity vector projected onto the local frame given a quaternion orientation.
 
     Args:
@@ -101,7 +101,7 @@ def get_projected_gravity_vector_from_quat(quat: jax.Array, eps: float = 1e-6) -
     return jnp.concatenate([gx, gy, -gz], axis=-1)
 
 
-def rotate_vector_by_quat(vector: jax.Array, quat: jax.Array, eps: float = 1e-6) -> jax.Array:
+def rotate_vector_by_quat(vector: Array, quat: Array, eps: float = 1e-6) -> Array:
     """Rotates a vector by a quaternion.
 
     Args:
@@ -156,3 +156,24 @@ def rotate_vector_by_quat(vector: jax.Array, quat: jax.Array, eps: float = 1e-6)
     )
 
     return jnp.concatenate([xx, yy, zz], axis=-1)
+
+
+def cubic_bezier_interpolation(y_start: Array, y_end: Array, x: Array) -> Array:
+    """Cubic bezier interpolation.
+
+    This is a cubic bezier curve that starts at y_start and ends at y_end,
+    and is controlled by the parameter x. The curve is defined by the following formula:
+
+    y(x) = y_start + (y_end - y_start) * (x**3 + 3 * (x**2 * (1 - x)))
+
+    Args:
+        y_start: The start value, shape (*).
+        y_end: The end value, shape (*).
+        x: The interpolation parameter, shape (*).
+
+    Returns:
+        The interpolated value, shape (*).
+    """
+    y_diff = y_end - y_start
+    bezier = x**3 + 3 * (x**2 * (1 - x))
+    return y_start + y_diff * bezier

xax/task/base.py

@@ -82,9 +82,6 @@ class BaseTask(Generic[Config]):
     def on_after_checkpoint_save(self, ckpt_path: Path, state: State) -> State:
         return state
 
-    def on_before_checkpoint_load(self, ckpt_path: Path) -> None:
-        pass
-
     @functools.cached_property
     def task_class_name(self) -> str:
         return self.__class__.__name__

xax/task/mixins/checkpointing.py

@@ -98,19 +98,39 @@ class CheckpointingMixin(ArtifactsMixin[Config], Generic[Config]):
     ) -> tuple[PyTree, State, DictConfig]: ...
 
     @overload
-    def load_checkpoint(self, path: Path, part: Literal["model"]) -> PyTree: ...
+    def load_checkpoint(
+        self,
+        path: Path,
+        part: Literal["model"],
+    ) -> PyTree: ...
 
     @overload
-    def load_checkpoint(self, path: Path, part: Literal["opt"]) -> optax.GradientTransformation: ...
+    def load_checkpoint(
+        self,
+        path: Path,
+        part: Literal["opt"],
+    ) -> optax.GradientTransformation: ...
 
     @overload
-    def load_checkpoint(self, path: Path, part: Literal["opt_state"]) -> optax.OptState: ...
+    def load_checkpoint(
+        self,
+        path: Path,
+        part: Literal["opt_state"],
+    ) -> optax.OptState: ...
 
     @overload
-    def load_checkpoint(self, path: Path, part: Literal["state"]) -> State: ...
+    def load_checkpoint(
+        self,
+        path: Path,
+        part: Literal["state"],
+    ) -> State: ...
 
     @overload
-    def load_checkpoint(self, path: Path, part: Literal["config"]) -> DictConfig: ...
+    def load_checkpoint(
+        self,
+        path: Path,
+        part: Literal["config"],
+    ) -> DictConfig: ...
 
     def load_checkpoint(
         self,
@@ -125,9 +145,6 @@ class CheckpointingMixin(ArtifactsMixin[Config], Generic[Config]):
         | State
         | DictConfig
     ):
-        # Calls the base callback.
-        self.on_before_checkpoint_load(path)
-
         with tarfile.open(path, "r:gz") as tar:
 
             def get_model() -> PyTree:

xax/task/mixins/compile.py

@@ -32,6 +32,10 @@ def get_cache_dir() -> str | None:
 @dataclass
 class CompileOptions:
     # JAX compilation options
+    debug_nans: bool = field(
+        value=False,
+        help="If True, breaks on NaNs",
+    )
     disable_jit: bool = field(
         value=False,
         help="If True, disables JIT compilation",
@@ -89,6 +93,10 @@ class CompileMixin(BaseTask[Config], Generic[Config]):
         cc = self.config.compile
 
         # Set basic compilation flags
+        if cc.debug_nans:
+            logger.info("Enabling NaNs debugging")
+            jax.config.update("jax_debug_nans", True)
+
         if cc.disable_jit:
             logger.info("Disabling JIT compilation")
             jax.config.update("jax_disable_jit", True)

xax/task/mixins/train.py

@@ -50,8 +50,7 @@ from xax.utils.experiments import (
     TrainingFinishedError,
     diff_configs,
     get_diff_string,
-    get_git_state,
-    get_packages_with_versions,
+    get_state_file_string,
     get_training_code,
 )
 from xax.utils.jax import jit as xax_jit
@@ -534,9 +533,8 @@ class TrainMixin(
         logger.log(LOG_STATUS, self.task_path)
         logger.log(LOG_STATUS, self.task_name)
         logger.log(LOG_STATUS, "JAX devices: %s", jax.devices())
-        self.logger.log_file("git_state.txt", get_git_state(self))
-        self.logger.log_file("packages.txt", get_packages_with_versions())
-        self.logger.log_file("training_code.txt", get_training_code(self))
+        self.logger.log_file("state.txt", get_state_file_string(self))
+        self.logger.log_file("training_code.py", get_training_code(self))
         self.logger.log_file("config.yaml", self.config_str(self.config, use_cli=False))
 
     def model_partition_fn(self, item: Any) -> bool:  # noqa: ANN401

xax/utils/debugging.py

@@ -53,3 +53,7 @@ def get_named_leaves(
 
 def breakpoint_if_nan(x: Array) -> None:
     jax.lax.cond(jnp.any(jnp.isnan(x)), lambda: jax.debug.breakpoint(), lambda: None)
+
+
+def log_if_nan(x: Array, loc: str) -> None:
+    jax.lax.cond(jnp.any(jnp.isnan(x)), lambda: jax.debug.print("=== NaNs: {loc} ===", loc=loc), lambda: None)

xax/utils/experiments.py

@@ -479,6 +479,20 @@ def get_packages_with_versions() -> str:
     return "\n".join([f"{key}=={version}" for key, version in sorted(packages)])
 
 
+def get_command_line_string() -> str:
+    return " ".join(sys.argv)
+
+
+def get_state_file_string(obj: object) -> str:
+    return "\n\n".join(
+        [
+            f"=== Command Line ===\n\n{get_command_line_string()}",
+            f"=== Git State ===\n\n{get_git_state(obj)}",
+            f"=== Packages ===\n\n{get_packages_with_versions()}",
+        ]
+    )
+
+
 def get_training_code(obj: object) -> str:
     """Gets the text from the file containing the provided object.
 

{xax-0.1.13.dist-info → xax-0.1.15.dist-info}/METADATA

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

{xax-0.1.13.dist-info → xax-0.1.15.dist-info}/RECORD

@@ -1,4 +1,4 @@
-xax/__init__.py,sha256=zBQiNKXLrE-9UPy1q-XBFwPNAKTAxr6wpAwYwaTVggs,13922
+xax/__init__.py,sha256=bV2mTcuiVaVNvwgbDgg7dKDkMeuyA0mqF0muU5KZHeg,14104
 xax/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/requirements-dev.txt,sha256=qkscNkFzWd1S5fump-AKH53rR65v2x5FmboFdy_kKvs,128
 xax/requirements.txt,sha256=9LAEZ5c5gqRSARRVA6xJsVTa4MebPZuC4yOkkwkZJFw,297
@@ -10,13 +10,13 @@ xax/nn/embeddings.py,sha256=bQGxBFxkLwi2MQLkRfGaHPH5P_KKB21HdI7VNWTKIOQ,11847
 xax/nn/equinox.py,sha256=5fdOKRXqAVZPsV-aEez3i1wamr_oBYnG74GP1jEthjM,4843
 xax/nn/export.py,sha256=7Yemw3T33QGEP8RkmTkpu6tRVOhut2RUJmttNFfCgFw,5537
 xax/nn/functions.py,sha256=CI_OmspaQwN9nl4hwefIU3_I7m6gBZwJ9aGK1JGUgr0,2713
-xax/nn/geom.py,sha256=Bj9Z4Y-uoNQuaA_eB_MyG7yImZLuOq8KCLUj1l3daoc,4545
+xax/nn/geom.py,sha256=PN0Ndn575aVtsSfxi67RghHB7luRkqtpS7bPbT1LpLE,5201
 xax/nn/losses.py,sha256=Q_NVnm5n4UPBvp5nI_1aUptfXnqFYoUeFwySiyvopHg,272
 xax/nn/norm.py,sha256=WgZ3QCrUnf-YecwhEtVPcr99fKK3ECl_UeiAs2uv7oo,564
 xax/nn/parallel.py,sha256=fnTiT7MsG7eQrJvqwjIz2Ifo3P27TuxIJzmpGYSa_dQ,4608
 xax/nn/ssm.py,sha256=8dLAcQ1hBaMT-kkHvwGu_ecxJeTY32WeMYmd4T4KtxA,10745
 xax/task/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-xax/task/base.py,sha256=E4l1yCrAkM2TVTbVYrmk6BoVHMkbD4IYsTT921XOyi0,7760
+xax/task/base.py,sha256=DqgGIlo5kEWpYix3DdPCEkCgVLUOocjyFr8okaSUq-k,7680
 xax/task/logger.py,sha256=1SZjVC6UCtZUoMPcpp3ckotL324QDeYDvHVhf5MHVqg,36271
 xax/task/script.py,sha256=bMMIJoUtpSBvPp6-7bejTrajTXvSg0794sYLKdPIToE,972
 xax/task/task.py,sha256=UHMpnv__gqMcfbC_L-Hhk-DCnUYlFVsgbNf-v8o8B7U,1424
@@ -32,8 +32,8 @@ xax/task/loggers/stdout.py,sha256=BBXqr95gNt5KuCN8XyKnTJF8JdwkR4JgLKrkvcaTBVM,67
 xax/task/loggers/tensorboard.py,sha256=kI8LvBuBBhPgkP8TeaTQb9SQ0FqaIodwQh2SuWDCnIA,7706
 xax/task/mixins/__init__.py,sha256=D3oU31rB9FeOr9MPLleLt5JFbftUr4sBTwgnwQdc2qA,809
 xax/task/mixins/artifacts.py,sha256=2ezmZGzPGe3nhsd9KRkeHWWXdbT9m7drzimIfw6v1XY,2892
-xax/task/mixins/checkpointing.py,sha256=a6tVyISsDIz68rrhb1rAh3rjQlqkDVJCmSBmETQrnRM,8480
-xax/task/mixins/compile.py,sha256=8jEdlGs-a14N_CwZA3Rxe461MT83dyIDr3Z56VkjviQ,3693
+xax/task/mixins/checkpointing.py,sha256=nRddgtasagf0oTZE9LE5IN5JY7jy4BD_M0rlqYp4sCM,8554
+xax/task/mixins/compile.py,sha256=PG5aF3W9v_xGiImHgUJ7gmwuQQoSQWufdpl2N_mlLX0,3922
 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
@@ -41,10 +41,10 @@ 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=aIebtOIvERYofSyqzNGBpNYlNrXweqFUqM9dHiTx3Dc,26253
+xax/task/mixins/train.py,sha256=1hmUx1HIL8HKfwOnupS3Knsw1CiK2YCbIQnUTYyDEms,26157
 xax/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-xax/utils/debugging.py,sha256=0DU41DvYp3SZ9tMrM7sSFpfhC7dieMYR7eRlGNAFrdM,1783
-xax/utils/experiments.py,sha256=5CUja1H_cx4dnVqTGQekOpIhqISwHtAgLxZ34GV7cwM,29229
+xax/utils/debugging.py,sha256=OtUdu-3tQsQtik0Q9UM-SNV46IbPjwrAfZcywzoB5d4,1940
+xax/utils/experiments.py,sha256=X6MESZ3z_Z0DLH6NQucuPzibuOc6rZmlf5UZt4in458,29591
 xax/utils/jax.py,sha256=tC0NNelbrSTzwNGluiwLGKtoHhVpgdzrv-xherB3VtY,4752
 xax/utils/jaxpr.py,sha256=S80nyEkv188RInzq3kCAdkQCU-bf6s0oPTrCE_LjkRs,2298
 xax/utils/logging.py,sha256=GAhTne2rdB4Fa1lzk06DMO15U8MTejn6XTClShC-ZtU,6622
@@ -58,8 +58,8 @@ xax/utils/data/collate.py,sha256=Rd9vMomr_S_zCa_Hi4dO-8ntzAfVwndIUtuXFA3iNcc,706
 xax/utils/types/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/utils/types/frozen_dict.py,sha256=ZCMGfSfr2_b2qZbq9ywPD0zej5tpVSId2JftXpwfB5k,4686
 xax/utils/types/hashable_array.py,sha256=l5iIcFmkYzfGeaZmcSoeFkthFASqM8xJYK3AXhZQYwc,992
-xax-0.1.13.dist-info/licenses/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
-xax-0.1.13.dist-info/METADATA,sha256=P3N5hJMZtXPs199OHPN_cBp57S9zlOGB2B1TRPaczuI,1878
-xax-0.1.13.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-xax-0.1.13.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
-xax-0.1.13.dist-info/RECORD,,
+xax-0.1.15.dist-info/licenses/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
+xax-0.1.15.dist-info/METADATA,sha256=i5thFSTL1Zx03UpnCj7f71rxSgs0P3L6ZDd6vYEtM7U,1878
+xax-0.1.15.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
+xax-0.1.15.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
+xax-0.1.15.dist-info/RECORD,,