xax 0.2.14__py3-none-any.whl → 0.2.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.2.14"
+__version__ = "0.2.15"
 
 # 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/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/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/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.dist-info → xax-0.2.15.dist-info}/METADATA

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

{xax-0.2.14.dist-info → xax-0.2.15.dist-info}/RECORD

@@ -1,4 +1,4 @@
-xax/__init__.py,sha256=9RtfWhU2Qb-YlZurazJ_GKpSpPaGs0nwu9-nmzbC0Pk,15510
+xax/__init__.py,sha256=JVxuGfbwBPHXiF4kSG0Pb73mzu3EIaRipjvt0Y-Z9W4,15733
 xax/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/requirements-dev.txt,sha256=qkscNkFzWd1S5fump-AKH53rR65v2x5FmboFdy_kKvs,128
 xax/requirements.txt,sha256=6qY-84e-sTmlfJNrSjwONQKqzAn5h8G_oGIhnhmfSr4,302
@@ -12,7 +12,7 @@ xax/nn/export.py,sha256=pRfM2B4hB2EvljysC6AjtgB_7Cn7JtaP3dhYU2stZtY,5545
 xax/nn/functions.py,sha256=bA5kJYzMtFM8eUqBC086i355zJMAO7k_vPFNSDBI9-s,2814
 xax/nn/geom.py,sha256=B8QE-L-xJWhf9KygTByPUAWe7Clpek4GlTABpsJFMBs,7702
 xax/nn/losses.py,sha256=Q_NVnm5n4UPBvp5nI_1aUptfXnqFYoUeFwySiyvopHg,272
-xax/nn/norm.py,sha256=WgZ3QCrUnf-YecwhEtVPcr99fKK3ECl_UeiAs2uv7oo,564
+xax/nn/metrics.py,sha256=OAkeScwhi-wTBIJ59KHUhYbZTq4V4V-LG-mKlxMJ7bY,3238
 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
@@ -41,7 +41,7 @@ xax/task/mixins/logger.py,sha256=6oXsJJyNUx6YT3q58FVXMZBUpMgjVkGre6BXFN20cVI,280
 xax/task/mixins/process.py,sha256=hqDEsMp_SL6ee97iq26-G0g49OcWZZaX82JD4F22eJU,1781
 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=Fgx2SWGC0e1QtRv1iTXXNg45dzbCzur3UHvjRZfOoiM,33465
+xax/task/mixins/train.py,sha256=sUgZ7_WI4GUreYIDSICpU81IFJNJiHlP0VSv3QFvAB4,33483
 xax/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/utils/debugging.py,sha256=OtUdu-3tQsQtik0Q9UM-SNV46IbPjwrAfZcywzoB5d4,1940
 xax/utils/experiments.py,sha256=bj8BftSHT3fFzfiJ0Co0WvqWo0rUS8kQnQYpVvH8FTM,29942
@@ -50,7 +50,7 @@ xax/utils/jaxpr.py,sha256=H7pWl48ROXIB1-ZPWYfOn-ou3EBMxYWIwc_A0reJQoo,2333
 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=Qp4u-jNOo9nd-2vdiuGHXFCTDERo4-zAGwJX_7kG7WM,9045
+xax/utils/pytree.py,sha256=rVY2kKa637xfX3Oue6OP9ScwmDyxJ_CeHkUpZZtmN04,9231
 xax/utils/tensorboard.py,sha256=P0oIFvX2Qts1H4lkpizhRIpQdD0MNppVMeut0Z94yCs,19878
 xax/utils/text.py,sha256=xS02aSzdywl3KIaNSpKWcxdd37oYlUJtu9wIjkc1wVc,10654
 xax/utils/data/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -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=ebtHENhyUzSjyJTlbMaLtcckQIJ7EtgJiok_40TJZpo,4689
 xax/utils/types/hashable_array.py,sha256=l5iIcFmkYzfGeaZmcSoeFkthFASqM8xJYK3AXhZQYwc,992
-xax-0.2.14.dist-info/licenses/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
-xax-0.2.14.dist-info/METADATA,sha256=5Agik9rI2VgDzlElHUVYpQm8JVOytLrcIxtNMCa_UUE,1880
-xax-0.2.14.dist-info/WHEEL,sha256=pxyMxgL8-pra_rKaQ4drOZAegBVuX-G_4nRHjjgWbmo,91
-xax-0.2.14.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
-xax-0.2.14.dist-info/RECORD,,
+xax-0.2.15.dist-info/licenses/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
+xax-0.2.15.dist-info/METADATA,sha256=6LJoiKOyNmF1MJSwVSdbEJATzSv1P77Amn4ZJCbWaP0,1880
+xax-0.2.15.dist-info/WHEEL,sha256=pxyMxgL8-pra_rKaQ4drOZAegBVuX-G_4nRHjjgWbmo,91
+xax-0.2.15.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
+xax-0.2.15.dist-info/RECORD,,

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}")