xax 0.2.14__tar.gz → 0.2.16__tar.gz

{xax-0.2.14/xax.egg-info → xax-0.2.16}/PKG-INFO

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

{xax-0.2.14 → xax-0.2.16}/xax/__init__.py

@@ -12,7 +12,7 @@ and running the update script:
     python -m scripts.update_api --inplace
 """
 
-__version__ = "0.2.14"
+__version__ = "0.2.16"
 
 # This list shouldn't be modified by hand; instead, run the update script.
 __all__ = [
@@ -51,6 +51,7 @@ __all__ = [
     "rotation_matrix_to_rotation6d",
     "cross_entropy",
     "cast_norm_type",
+    "dynamic_time_warping",
     "get_norm",
     "is_master",
     "BaseSSMBlock",
@@ -136,6 +137,7 @@ __all__ = [
     "reshuffle_pytree_independently",
     "slice_array",
     "slice_pytree",
+    "tuple_insert",
     "update_pytree",
     "TextBlock",
     "camelcase_to_snakecase",
@@ -229,8 +231,9 @@ NAME_MAP: dict[str, str] = {
     "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",
+    "cast_norm_type": "nn.metrics",
+    "dynamic_time_warping": "nn.metrics",
+    "get_norm": "nn.metrics",
     "is_master": "nn.parallel",
     "BaseSSMBlock": "nn.ssm",
     "DiagSSMBlock": "nn.ssm",
@@ -315,6 +318,7 @@ NAME_MAP: dict[str, str] = {
     "reshuffle_pytree_independently": "utils.pytree",
     "slice_array": "utils.pytree",
     "slice_pytree": "utils.pytree",
+    "tuple_insert": "utils.pytree",
     "update_pytree": "utils.pytree",
     "TextBlock": "utils.text",
     "camelcase_to_snakecase": "utils.text",
@@ -345,7 +349,7 @@ NAME_MAP.update(
         "LOG_ERROR_SUMMARY": "utils.logging",
         "LOG_PING": "utils.logging",
         "LOG_STATUS": "utils.logging",
-        "NormType": "nn.norm",
+        "NormType": "nn.metrics",
         "Output": "task.mixins.output",
         "Phase": "core.state",
         "RawConfigType": "task.base",
@@ -410,7 +414,12 @@ if IMPORT_ALL or TYPE_CHECKING:
         rotation_matrix_to_rotation6d,
     )
     from xax.nn.losses import cross_entropy
-    from xax.nn.norm import NormType, cast_norm_type, get_norm
+    from xax.nn.metrics import (
+        NormType,
+        cast_norm_type,
+        dynamic_time_warping,
+        get_norm,
+    )
     from xax.nn.parallel import is_master
     from xax.nn.ssm import SSM, BaseSSMBlock, DiagSSMBlock, SSMBlock
     from xax.task.base import RawConfigType
@@ -495,6 +504,7 @@ if IMPORT_ALL or TYPE_CHECKING:
         reshuffle_pytree_independently,
         slice_array,
         slice_pytree,
+        tuple_insert,
         update_pytree,
     )
     from xax.utils.text import (

{xax-0.2.14 → xax-0.2.16}/xax/nn/geom.py

@@ -102,12 +102,13 @@ def get_projected_gravity_vector_from_quat(quat: Array, eps: float = 1e-6) -> Ar
     return jnp.concatenate([gx, gy, -gz], axis=-1)
 
 
-def rotate_vector_by_quat(vector: Array, quat: Array, eps: float = 1e-6) -> Array:
+def rotate_vector_by_quat(vector: Array, quat: Array, inverse: bool = False, eps: float = 1e-6) -> Array:
     """Rotates a vector by a quaternion.
 
     Args:
         vector: The vector to rotate, shape (*, 3).
         quat: The quaternion to rotate by, shape (*, 4).
+        inverse: If True, rotate the vector by the conjugate of the quaternion.
         eps: A small epsilon value to avoid division by zero.
 
     Returns:
@@ -117,6 +118,9 @@ def rotate_vector_by_quat(vector: Array, quat: Array, eps: float = 1e-6) -> Arra
     quat = quat / (jnp.linalg.norm(quat, axis=-1, keepdims=True) + eps)
     w, x, y, z = jnp.split(quat, 4, axis=-1)
 
+    if inverse:
+        x, y, z = -x, -y, -z
+
     # Extract vector components
     vx, vy, vz = jnp.split(vector, 3, axis=-1)
 

xax-0.2.16/xax/nn/metrics.py

@@ -0,0 +1,92 @@
+"""Norm and metric utilities."""
+
+from typing import Literal, cast, get_args, overload
+
+import chex
+import jax
+import jax.numpy as jnp
+from jaxtyping import Array
+
+from xax.utils.jax import jit as xax_jit
+
+NormType = Literal["l1", "l2"]
+
+
+def cast_norm_type(norm: str) -> NormType:
+    if norm not in get_args(NormType):
+        raise ValueError(f"Invalid norm: {norm}")
+    return cast(NormType, norm)
+
+
+def get_norm(x: Array, norm: NormType) -> Array:
+    match norm:
+        case "l1":
+            return jnp.abs(x)
+        case "l2":
+            return jnp.square(x)
+        case _:
+            raise ValueError(f"Invalid norm: {norm}")
+
+
+@overload
+def dynamic_time_warping(distance_matrix_nm: Array) -> Array: ...
+
+
+@overload
+def dynamic_time_warping(distance_matrix_nm: Array, return_path: Literal[True]) -> tuple[Array, Array]: ...
+
+
+@xax_jit(static_argnames=["return_path"])
+def dynamic_time_warping(distance_matrix_nm: Array, return_path: bool = False) -> Array | tuple[Array, Array]:
+    """Dynamic Time Warping.
+
+    Args:
+        distance_matrix_nm: A matrix of pairwise distances between two
+            sequences, with shape (N, M), with the condition that N <= M.
+        return_path: If set, return the minimum path, otherwise just return
+            the cost. The latter is preferred if using this function as a
+            distance metric since it avoids the backwards scan on backpointers.
+
+    Returns:
+        The cost of the minimum path from the top-left corner of the distance
+        matrix to the bottom-right corner, along with the indices of that
+        minimum path.
+    """
+    chex.assert_shape(distance_matrix_nm, (None, None))
+    n, m = distance_matrix_nm.shape
+
+    assert n <= m, f"Invalid dynamic time warping distance matrix shape: ({n}, {m})"
+
+    # Masks values which cannot be reached.
+    row_idx = jnp.arange(n)[:, None]
+    col_idx = jnp.arange(m)[None, :]
+    mask = row_idx > col_idx
+    distance_matrix_nm = jnp.where(mask, jnp.inf, distance_matrix_nm)
+
+    # Pre-pads with inf
+    distance_matrix_nm = jnp.pad(distance_matrix_nm, ((1, 0), (0, 0)), mode="constant", constant_values=jnp.inf)
+    indices = jnp.arange(n)
+
+    # Scan over remaining rows to fill cost matrix
+    def scan_fn(prev_cost: Array, cur_distances: Array) -> tuple[Array, Array]:
+        same_trans = prev_cost
+        prev_trans = jnp.pad(prev_cost[:-1], ((1, 0),), mode="constant", constant_values=jnp.inf)
+        nc = jnp.minimum(prev_trans, same_trans) + cur_distances[1:]
+        return nc, jnp.where(prev_trans < same_trans, indices - 1, indices) if return_path else nc
+
+    init_cost = distance_matrix_nm[1:, 0]
+    final_cost, back_pointers = jax.lax.scan(scan_fn, init_cost, distance_matrix_nm[:, 1:].T)
+
+    if not return_path:
+        return final_cost
+
+    # Scan the back pointers backwards to get the minimum path.
+    def scan_back_fn(carry: Array, back_pointer: Array) -> tuple[Array, Array]:
+        prev_idx = back_pointer[carry]
+        return prev_idx, carry
+
+    final_index = jnp.array(n - 1)
+    _, min_path = jax.lax.scan(scan_back_fn, final_index, back_pointers, reverse=True)
+    min_path = jnp.pad(min_path, ((1, 0)), mode="constant", constant_values=0)
+
+    return final_cost[-1], min_path

{xax-0.2.14 → xax-0.2.16}/xax/task/mixins/train.py

@@ -363,7 +363,7 @@ class TrainMixin(
         self,
         key: PRNGKeyArray,
         load_optimizer: Literal[True],
-    ) -> tuple[PyTree, optax.GradientTransformation, optax.OptState, State]: ...
+    ) -> tuple[list[PyTree], list[optax.GradientTransformation], list[optax.OptState], State]: ...
 
     def load_initial_state(
         self,

{xax-0.2.14 → xax-0.2.16}/xax/utils/pytree.py

@@ -1,5 +1,7 @@
 """Utils for accessing, modifying, and otherwise manipulating pytrees."""
 
+from typing import TypeVar
+
 import chex
 import equinox as eqx
 import jax
@@ -7,6 +9,8 @@ import jax.numpy as jnp
 from jax import Array
 from jaxtyping import PRNGKeyArray, PyTree
 
+T = TypeVar("T")
+
 
 def slice_array(x: Array, start: Array, slice_length: int) -> Array:
     """Get a slice of an array along the first dimension.
@@ -243,3 +247,9 @@ 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))
+
+
+def tuple_insert(t: tuple[T, ...], index: int, value: T) -> tuple[T, ...]:
+    mut = list(t)
+    mut[index] = value
+    return tuple(mut)

{xax-0.2.14 → xax-0.2.16/xax.egg-info}/PKG-INFO

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

{xax-0.2.14 → xax-0.2.16}/xax.egg-info/SOURCES.txt

@@ -23,7 +23,7 @@ xax/nn/export.py
 xax/nn/functions.py
 xax/nn/geom.py
 xax/nn/losses.py
-xax/nn/norm.py
+xax/nn/metrics.py
 xax/nn/parallel.py
 xax/nn/ssm.py
 xax/task/__init__.py

xax-0.2.14/xax/nn/norm.py

@@ -1,24 +0,0 @@
-"""Normalization utilities."""
-
-from typing import Literal, cast, get_args
-
-import jax.numpy as jnp
-from jaxtyping import Array
-
-NormType = Literal["l1", "l2"]
-
-
-def cast_norm_type(norm: str) -> NormType:
-    if norm not in get_args(NormType):
-        raise ValueError(f"Invalid norm: {norm}")
-    return cast(NormType, norm)
-
-
-def get_norm(x: Array, norm: NormType) -> Array:
-    match norm:
-        case "l1":
-            return jnp.abs(x)
-        case "l2":
-            return jnp.square(x)
-        case _:
-            raise ValueError(f"Invalid norm: {norm}")