xax 0.0.6__py3-none-any.whl → 0.0.7__py3-none-any.whl

xax/__init__.py

@@ -11,7 +11,7 @@ This file can be maintained by running the update script:
     python -m scripts.update_api --inplace
 """
 
-__version__ = "0.0.6"
+__version__ = "0.0.7"
 
 # This list shouldn't be modified by hand; instead, run the update script.
 __all__ = [
@@ -34,6 +34,8 @@ __all__ = [
     "get_positional_embeddings",
     "get_rotary_embeddings",
     "rotary_embeddings",
+    "euler_to_quat",
+    "quat_to_euler",
     "is_master",
     "BaseLauncher",
     "CliLauncher",
@@ -78,11 +80,18 @@ __all__ = [
     "save_config",
     "stage_environment",
     "to_markdown_table",
+    "jit",
     "ColoredFormatter",
     "configure_logging",
     "one_hot",
     "partial_flatten",
     "worker_chunk",
+    "profile",
+    "compute_nan_ratio",
+    "flatten_array",
+    "flatten_pytree",
+    "slice_array",
+    "slice_pytree",
     "TextBlock",
     "camelcase_to_snakecase",
     "colored",
@@ -142,6 +151,8 @@ NAME_MAP: dict[str, str] = {
     "get_positional_embeddings": "nn.embeddings",
     "get_rotary_embeddings": "nn.embeddings",
     "rotary_embeddings": "nn.embeddings",
+    "euler_to_quat": "nn.geom",
+    "quat_to_euler": "nn.geom",
     "is_master": "nn.parallel",
     "BaseLauncher": "task.launchers.base",
     "CliLauncher": "task.launchers.cli",
@@ -186,11 +197,18 @@ NAME_MAP: dict[str, str] = {
     "save_config": "utils.experiments",
     "stage_environment": "utils.experiments",
     "to_markdown_table": "utils.experiments",
+    "jit": "utils.jax",
     "ColoredFormatter": "utils.logging",
     "configure_logging": "utils.logging",
     "one_hot": "utils.numpy",
     "partial_flatten": "utils.numpy",
     "worker_chunk": "utils.numpy",
+    "profile": "utils.profile",
+    "compute_nan_ratio": "utils.pytree",
+    "flatten_array": "utils.pytree",
+    "flatten_pytree": "utils.pytree",
+    "slice_array": "utils.pytree",
+    "slice_pytree": "utils.pytree",
     "TextBlock": "utils.text",
     "camelcase_to_snakecase": "utils.text",
     "colored": "utils.text",
@@ -257,6 +275,7 @@ if IMPORT_ALL or TYPE_CHECKING:
         get_rotary_embeddings,
         rotary_embeddings,
     )
+    from xax.nn.geom import euler_to_quat, quat_to_euler
     from xax.nn.parallel import is_master
     from xax.task.base import RawConfigType
     from xax.task.launchers.base import BaseLauncher
@@ -299,6 +318,7 @@ if IMPORT_ALL or TYPE_CHECKING:
         stage_environment,
         to_markdown_table,
     )
+    from xax.utils.jax import jit
     from xax.utils.logging import (
         LOG_ERROR_SUMMARY,
         LOG_PING,
@@ -307,6 +327,14 @@ if IMPORT_ALL or TYPE_CHECKING:
         configure_logging,
     )
     from xax.utils.numpy import one_hot, partial_flatten, worker_chunk
+    from xax.utils.profile import profile
+    from xax.utils.pytree import (
+        compute_nan_ratio,
+        flatten_array,
+        flatten_pytree,
+        slice_array,
+        slice_pytree,
+    )
     from xax.utils.text import (
         TextBlock,
         camelcase_to_snakecase,

xax/nn/geom.py

@@ -0,0 +1,75 @@
+"""Defines geometry functions."""
+
+import jax
+from jax import numpy as jnp
+
+
+def quat_to_euler(quat_4: jax.Array, eps: float = 1e-6) -> jax.Array:
+    """Normalizes and converts a quaternion (w, x, y, z) to roll, pitch, yaw.
+
+    Args:
+        quat_4: The quaternion to convert, shape (*, 4).
+        eps: A small epsilon value to avoid division by zero.
+
+    Returns:
+        The roll, pitch, yaw angles with shape (*, 3).
+    """
+    quat_4 = quat_4 / (jnp.linalg.norm(quat_4, axis=-1, keepdims=True) + eps)
+    w, x, y, z = jnp.split(quat_4, 4, axis=-1)
+
+    # Roll (x-axis rotation)
+    sinr_cosp = 2.0 * (w * x + y * z)
+    cosr_cosp = 1.0 - 2.0 * (x * x + y * y)
+    roll = jnp.arctan2(sinr_cosp, cosr_cosp)
+
+    # Pitch (y-axis rotation)
+    sinp = 2.0 * (w * y - z * x)
+
+    # Handle edge cases where |sinp| >= 1
+    pitch = jnp.where(
+        jnp.abs(sinp) >= 1.0,
+        jnp.sign(sinp) * jnp.pi / 2.0,  # Use 90 degrees if out of range
+        jnp.arcsin(sinp),
+    )
+
+    # Yaw (z-axis rotation)
+    siny_cosp = 2.0 * (w * z + x * y)
+    cosy_cosp = 1.0 - 2.0 * (y * y + z * z)
+    yaw = jnp.arctan2(siny_cosp, cosy_cosp)
+
+    return jnp.concatenate([roll, pitch, yaw], axis=-1)
+
+
+def euler_to_quat(euler_3: jax.Array) -> jax.Array:
+    """Converts roll, pitch, yaw angles to a quaternion (w, x, y, z).
+
+    Args:
+        euler_3: The roll, pitch, yaw angles, shape (*, 3).
+
+    Returns:
+        The quaternion with shape (*, 4).
+    """
+    # Extract roll, pitch, yaw from input
+    roll, pitch, yaw = jnp.split(euler_3, 3, axis=-1)
+
+    # Calculate trigonometric functions for each angle
+    cr = jnp.cos(roll * 0.5)
+    sr = jnp.sin(roll * 0.5)
+    cp = jnp.cos(pitch * 0.5)
+    sp = jnp.sin(pitch * 0.5)
+    cy = jnp.cos(yaw * 0.5)
+    sy = jnp.sin(yaw * 0.5)
+
+    # Calculate quaternion components using the conversion formula
+    w = cr * cp * cy + sr * sp * sy
+    x = sr * cp * cy - cr * sp * sy
+    y = cr * sp * cy + sr * cp * sy
+    z = cr * cp * sy - sr * sp * cy
+
+    # Combine into quaternion [w, x, y, z]
+    quat = jnp.concatenate([w, x, y, z], axis=-1)
+
+    # Normalize the quaternion
+    quat = quat / jnp.linalg.norm(quat, axis=-1, keepdims=True)
+
+    return quat

xax/utils/jax.py

@@ -1,14 +1,140 @@
 """Defines some utility functions for interfacing with Jax."""
 
+import inspect
+import logging
+import os
+import time
+from functools import wraps
+from typing import Any, Callable, Iterable, ParamSpec, Sequence, TypeVar, cast
+
+import jax
 import jax.numpy as jnp
 import numpy as np
+from jax._src import sharding_impls
+from jax._src.lib import xla_client as xc
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_COMPILE_TIMEOUT = 1.0
 
 Number = int | float | np.ndarray | jnp.ndarray
 
 
+P = ParamSpec("P")  # For function parameters
+R = TypeVar("R")  # For function return type
+
+
 def as_float(value: int | float | np.ndarray | jnp.ndarray) -> float:
     if isinstance(value, (int, float)):
         return float(value)
     if isinstance(value, (np.ndarray, jnp.ndarray)):
         return float(value.item())
     raise TypeError(f"Unexpected type: {type(value)}")
+
+
+def get_hash(obj: object) -> int:
+    """Get a hash of an object.
+
+    If the object is hashable, use the hash. Otherwise, use the id.
+    """
+    if hasattr(obj, "__hash__"):
+        return hash(obj)
+    return id(obj)
+
+
+def jit(
+    in_shardings: Any = sharding_impls.UNSPECIFIED,  # noqa: ANN401
+    out_shardings: Any = sharding_impls.UNSPECIFIED,  # noqa: ANN401
+    static_argnums: int | Sequence[int] | None = None,
+    static_argnames: str | Iterable[str] | None = None,
+    donate_argnums: int | Sequence[int] | None = None,
+    donate_argnames: str | Iterable[str] | None = None,
+    keep_unused: bool = False,
+    device: xc.Device | None = None,
+    backend: str | None = None,
+    inline: bool = False,
+    abstracted_axes: Any | None = None,  # noqa: ANN401
+    compiler_options: dict[str, Any] | None = None,
+) -> Callable[[Callable[P, R]], Callable[P, R]]:
+    """Wrapper function that provides utility improvements over Jax's JIT.
+
+    Specifically, this function works on class methods, is toggleable, and
+    detects recompilations by matching hash values.
+
+    This is meant to be used as a decorator factory, and the decorated function
+    calls `wrapped`.
+    """
+
+    def decorator(fn: Callable[P, R]) -> Callable[P, R]:
+        class JitState:
+            compilation_count = 0
+            last_arg_dict: dict[str, int] | None = None
+
+        sig = inspect.signature(fn)
+        param_names = list(sig.parameters.keys())
+
+        jitted_fn = jax.jit(
+            fn,
+            in_shardings=in_shardings,
+            out_shardings=out_shardings,
+            static_argnums=static_argnums,
+            static_argnames=static_argnames,
+            donate_argnums=donate_argnums,
+            donate_argnames=donate_argnames,
+            keep_unused=keep_unused,
+            device=device,
+            backend=backend,
+            inline=inline,
+            abstracted_axes=abstracted_axes,
+            compiler_options=compiler_options,
+        )
+
+        @wraps(fn)
+        def wrapped(*args: P.args, **kwargs: P.kwargs) -> R:
+            if os.environ.get("DEBUG", "0") == "1":  # skipping during debug
+                return fn(*args, **kwargs)
+
+            do_profile = os.environ.get("JIT_PROFILE", "0") == "1"
+
+            if do_profile:
+                class_name = (args[0].__class__.__name__) + "." if fn.__name__ == "__call__" else ""
+                logger.info(
+                    "Currently running %s (count: %s)",
+                    f"{class_name}{fn.__name__}",
+                    JitState.compilation_count,
+                )
+
+            start_time = time.time()
+            res = jitted_fn(*args, **kwargs)
+            end_time = time.time()
+            runtime = end_time - start_time
+
+            # if this is true, if runtime is higher than COMPILE_TIMEOUT, we recompile
+            # TODO: we should probably reimplement the lower-level jitting logic to avoid this
+            if do_profile:
+                arg_dict = {}
+                for i, arg in enumerate(args):
+                    if i < len(param_names):
+                        arg_dict[param_names[i]] = get_hash(arg)
+                for k, v in kwargs.items():
+                    arg_dict[k] = get_hash(v)
+
+                logger.info("Hashing took %s seconds", runtime)
+                JitState.compilation_count += 1
+
+                if JitState.last_arg_dict is not None:
+                    all_keys = set(arg_dict.keys()) | set(JitState.last_arg_dict.keys())
+                    for k in all_keys:
+                        prev = JitState.last_arg_dict.get(k, "N/A")
+                        curr = arg_dict.get(k, "N/A")
+
+                        if prev != curr:
+                            logger.info("- Arg '%s' hash changed: %s -> %s", k, prev, curr)
+
+                JitState.last_arg_dict = arg_dict
+
+            return cast(R, res)
+
+        return wrapped
+
+    return decorator

xax/utils/profile.py

@@ -0,0 +1,61 @@
+"""Profiling utilities."""
+
+import logging
+import os
+import time
+from functools import wraps
+from typing import Callable, ParamSpec, TypeVar
+
+logger = logging.getLogger(__name__)
+
+P = ParamSpec("P")  # For function parameters
+R = TypeVar("R")  # For function return type
+
+
+def profile(fn: Callable[P, R]) -> Callable[P, R]:
+    """Profiling decorator that tracks function call count and execution time.
+
+    Activated when the PROFILE environment variable is set to "1".
+
+    Returns:
+        A decorated function with profiling capabilities.
+    """
+
+    class ProfileState:
+        call_count = 0
+        total_time = 0.0
+
+    @wraps(fn)
+    def wrapped(*args: P.args, **kwargs: P.kwargs) -> R:
+        if os.environ.get("PROFILE", "0") != "1":
+            return fn(*args, **kwargs)
+
+        start_time = time.time()
+        res = fn(*args, **kwargs)
+        end_time = time.time()
+        runtime = end_time - start_time
+
+        ProfileState.call_count += 1
+        ProfileState.total_time += runtime
+
+        # Handle class methods by showing class name
+        if fn.__name__ == "__call__" or (args and hasattr(args[0], "__class__")):
+            try:
+                class_name = args[0].__class__.__name__ + "."
+            except (IndexError, AttributeError):
+                class_name = ""
+        else:
+            class_name = ""
+
+        logger.info(
+            "%s %s - call #%s, took %s seconds, total: %s seconds",
+            class_name,
+            fn.__name__,
+            ProfileState.call_count,
+            runtime,
+            ProfileState.total_time,
+        )
+
+        return res
+
+    return wrapped

xax/utils/pytree.py

@@ -0,0 +1,50 @@
+"""Utils for accessing, modifying, and otherwise manipulating pytrees."""
+
+import chex
+import jax
+import jax.numpy as jnp
+from jax import Array
+from jaxtyping import PyTree
+
+
+def slice_array(x: Array, start: Array, slice_length: int) -> Array:
+    """Get a slice of an array along the first dimension.
+
+    For multi-dimensional arrays, this slices only along the first dimension
+    and keeps all other dimensions intact.
+    """
+    chex.assert_shape(start, ())
+    chex.assert_shape(slice_length, ())
+    start_indices = (start,) + (0,) * (len(x.shape) - 1)
+    slice_sizes = (slice_length,) + x.shape[1:]
+
+    return jax.lax.dynamic_slice(x, start_indices, slice_sizes)
+
+
+def slice_pytree(pytree: PyTree, start: Array, slice_length: int) -> PyTree:
+    """Get a slice of a pytree."""
+    return jax.tree_util.tree_map(lambda x: slice_array(x, start, slice_length), pytree)
+
+
+def flatten_array(x: Array, flatten_size: int) -> Array:
+    """Flatten an array into a (flatten_size, ...) array."""
+    reshaped = jnp.reshape(x, (flatten_size, *x.shape[2:]))
+    assert reshaped.shape[0] == flatten_size
+    return reshaped
+
+
+def flatten_pytree(pytree: PyTree, flatten_size: int) -> PyTree:
+    """Flatten a pytree into a (flatten_size, ...) pytree."""
+    return jax.tree_util.tree_map(lambda x: flatten_array(x, flatten_size), pytree)
+
+
+def compute_nan_ratio(pytree: PyTree) -> Array:
+    """Computes the ratio of NaNs vs non-NaNs in a given PyTree."""
+    nan_counts = jax.tree_util.tree_map(lambda x: jnp.sum(jnp.isnan(x)), pytree)
+    total_counts = jax.tree_util.tree_map(lambda x: x.size, pytree)
+
+    total_nans = jax.tree_util.tree_reduce(lambda a, b: a + b, nan_counts, 0)
+    total_elements = jax.tree_util.tree_reduce(lambda a, b: a + b, total_counts, 0)
+    overall_nan_ratio = jnp.array(total_nans / total_elements)
+
+    return overall_nan_ratio

{xax-0.0.6.dist-info → xax-0.0.7.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: xax
-Version: 0.0.6
+Version: 0.0.7
 Summary: The xax project
 Home-page: https://github.com/dpshai/xax
 Author: Benjamin Bolte

{xax-0.0.6.dist-info → xax-0.0.7.dist-info}/RECORD

@@ -1,4 +1,4 @@
-xax/__init__.py,sha256=RTUsDh_R0TFa09q-_U0vd-eCYRC-bCaHqHlayp8U2hU,9736
+xax/__init__.py,sha256=ScTkvKaxgpuKhhs9RINJa2XWCj899ndSYrB3FtScfxw,10509
 xax/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/requirements-dev.txt,sha256=qkscNkFzWd1S5fump-AKH53rR65v2x5FmboFdy_kKvs,128
 xax/requirements.txt,sha256=NmU9PNJhfLtNqqtWWf8WqMjgbBPCn_yt8oMGAgS7Fno,291
@@ -8,6 +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/functions.py,sha256=CI_OmspaQwN9nl4hwefIU3_I7m6gBZwJ9aGK1JGUgr0,2713
+xax/nn/geom.py,sha256=MtVar9AdqrJQGIFxcIFHyFnV_fblf9Pc4kQT_gTQASI,2195
 xax/nn/parallel.py,sha256=fnTiT7MsG7eQrJvqwjIz2Ifo3P27TuxIJzmpGYSa_dQ,4608
 xax/task/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/task/base.py,sha256=LHDmM2c_Ps5cGEzn_QUpmyInD7zJJm3Yt9eSeij2Vus,7297
@@ -38,15 +39,17 @@ xax/task/mixins/step_wrapper.py,sha256=DJw42mUGwgKx2tkeqatKR9_F4J8ug4wmxKMeJPmhc
 xax/task/mixins/train.py,sha256=dhGL_IuDaJy39BooYlO7JO-_EotKldtBhBplDGU_AnM,21745
 xax/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/utils/experiments.py,sha256=qT3H0fyVH8DN417x7T0Xmz4SKoogW81-EHcZfyktFI8,28300
-xax/utils/jax.py,sha256=VzEVB766UyH3_cgN6UP0FkCsDuGlYg5KJj8YJS4yYUk,439
+xax/utils/jax.py,sha256=tC0NNelbrSTzwNGluiwLGKtoHhVpgdzrv-xherB3VtY,4752
 xax/utils/logging.py,sha256=ST1hp2C2xntVVJBUHwo3YxPK19fBLNvHU2WGO1xqcXA,6418
 xax/utils/numpy.py,sha256=_jOXVi-d2AtJnRftPkRK5MDMzsU8slgw-Jjv4GRm6ns,1197
+xax/utils/profile.py,sha256=-aFdWpgYFvBsBZXSLL4zXrFe3zzsDqzmx4q5f2WOtpQ,1628
+xax/utils/pytree.py,sha256=Jwx6ErJfv1r2D23D4eKz1Hoo3mAJ0SEqC3EagZarWkw,1858
 xax/utils/tensorboard.py,sha256=oGq2E3Yr0z2xaACv2UOVt_CHEVc8fBxI8V1M99Fd34E,9742
 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.0.6.dist-info/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
-xax-0.0.6.dist-info/METADATA,sha256=YO2c2PUMWkH1ILfPhFWKK4Sodbo9qUpUOCIkm4aLHfg,1171
-xax-0.0.6.dist-info/WHEEL,sha256=jB7zZ3N9hIM9adW7qlTAyycLYW9npaWKLRzaoVcLKcM,91
-xax-0.0.6.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
-xax-0.0.6.dist-info/RECORD,,
+xax-0.0.7.dist-info/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
+xax-0.0.7.dist-info/METADATA,sha256=hE0KO4kYcN6Ed8iZ4649R5ENOUaQysBMW9vTh-94d4I,1171
+xax-0.0.7.dist-info/WHEEL,sha256=52BFRY2Up02UkjOa29eZOS2VxUrpPORXg1pkohGGUS8,91
+xax-0.0.7.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
+xax-0.0.7.dist-info/RECORD,,

{xax-0.0.6.dist-info → xax-0.0.7.dist-info}/WHEEL

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