xax 0.1.14__tar.gz → 0.1.15__tar.gz

{xax-0.1.14/xax.egg-info → xax-0.1.15}/PKG-INFO

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

{xax-0.1.14 → xax-0.1.15}/xax/__init__.py

@@ -12,7 +12,7 @@ and running the update script:
     python -m scripts.update_api --inplace
 """
 
-__version__ = "0.1.14"
+__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",
@@ -201,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",
@@ -363,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,

{xax-0.1.14 → xax-0.1.15}/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-0.1.14 → xax-0.1.15}/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-0.1.14 → xax-0.1.15}/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.14 → xax-0.1.15/xax.egg-info}/PKG-INFO

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