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

xax/__init__.py

@@ -4,14 +4,15 @@ This package is structured so that all the important stuff can be accessed
 without having to dig around through the internals. This is done by lazily
 importing the module by name.
 
-This file can be maintained by running the update script:
+This file can be maintained by updating the imports at the bottom of the file
+and running the update script:
 
 .. code-block:: bash
 
     python -m scripts.update_api --inplace
 """
 
-__version__ = "0.0.7"
+__version__ = "0.1.0"
 
 # This list shouldn't be modified by hand; instead, run the update script.
 __all__ = [
@@ -34,8 +35,20 @@ __all__ = [
     "get_positional_embeddings",
     "get_rotary_embeddings",
     "rotary_embeddings",
+    "MLPHyperParams",
+    "export_eqx_mlp",
+    "load_eqx",
+    "load_eqx_mlp",
+    "make_eqx_mlp",
+    "save_eqx",
+    "export",
+    "export_flax",
+    "export_with_params",
     "euler_to_quat",
+    "get_projected_gravity_vector_from_quat",
     "quat_to_euler",
+    "cast_norm_type",
+    "get_norm",
     "is_master",
     "BaseLauncher",
     "CliLauncher",
@@ -52,13 +65,16 @@ __all__ = [
     "CPUStatsOptions",
     "DataloaderConfig",
     "GPUStatsOptions",
+    "StepContext",
     "Script",
     "ScriptConfig",
     "Config",
     "Task",
     "collate",
     "collate_non_null",
+    "get_named_leaves",
     "BaseFileDownloader",
+    "ContextTimer",
     "CumulativeTimer",
     "DataDownloader",
     "IntervalTicker",
@@ -81,6 +97,7 @@ __all__ = [
     "stage_environment",
     "to_markdown_table",
     "jit",
+    "save_jaxpr_dot",
     "ColoredFormatter",
     "configure_logging",
     "one_hot",
@@ -90,8 +107,13 @@ __all__ = [
     "compute_nan_ratio",
     "flatten_array",
     "flatten_pytree",
+    "pytree_has_nans",
+    "reshuffle_pytree",
+    "reshuffle_pytree_along_dims",
+    "reshuffle_pytree_independently",
     "slice_array",
     "slice_pytree",
+    "update_pytree",
     "TextBlock",
     "camelcase_to_snakecase",
     "colored",
@@ -113,21 +135,36 @@ __all__ += [
     "Batch",
     "CollateMode",
     "EmbeddingKind",
+    "ActivationFunction",
+    "DTYPE",
     "LOG_ERROR_SUMMARY",
     "LOG_PING",
     "LOG_STATUS",
+    "NormType",
     "Output",
     "Phase",
     "RawConfigType",
 ]
 
 import os
+import shutil
 from typing import TYPE_CHECKING
 
+# Sets some useful XLA flags.
+xla_flags: list[str] = []
+if "XLA_FLAGS" in os.environ:
+    xla_flags.append(os.environ["XLA_FLAGS"])
+
+# If Nvidia GPU is detected (meaning, is `nvidia-smi` available?), disable
+# Triton GEMM kernels. See https://github.com/NVIDIA/JAX-Toolbox
+if shutil.which("nvidia-smi") is not None:
+    xla_flags += ["--xla_gpu_enable_latency_hiding_scheduler", "--xla_gpu_enable_triton_gemm"]
+os.environ["XLA_FLAGS"] = " ".join(xla_flags)
+
 # If this flag is set, eagerly imports the entire package (not recommended).
 IMPORT_ALL = int(os.environ.get("XAX_IMPORT_ALL", "0")) != 0
 
-del os
+del os, shutil, xla_flags
 
 # This dictionary is auto-generated and shouldn't be modified by hand; instead,
 # run the update script.
@@ -151,8 +188,20 @@ NAME_MAP: dict[str, str] = {
     "get_positional_embeddings": "nn.embeddings",
     "get_rotary_embeddings": "nn.embeddings",
     "rotary_embeddings": "nn.embeddings",
+    "MLPHyperParams": "nn.equinox",
+    "export_eqx_mlp": "nn.equinox",
+    "load_eqx": "nn.equinox",
+    "load_eqx_mlp": "nn.equinox",
+    "make_eqx_mlp": "nn.equinox",
+    "save_eqx": "nn.equinox",
+    "export": "nn.export",
+    "export_flax": "nn.export",
+    "export_with_params": "nn.export",
     "euler_to_quat": "nn.geom",
+    "get_projected_gravity_vector_from_quat": "nn.geom",
     "quat_to_euler": "nn.geom",
+    "cast_norm_type": "nn.norm",
+    "get_norm": "nn.norm",
     "is_master": "nn.parallel",
     "BaseLauncher": "task.launchers.base",
     "CliLauncher": "task.launchers.cli",
@@ -169,13 +218,16 @@ NAME_MAP: dict[str, str] = {
     "CPUStatsOptions": "task.mixins.cpu_stats",
     "DataloaderConfig": "task.mixins.data_loader",
     "GPUStatsOptions": "task.mixins.gpu_stats",
+    "StepContext": "task.mixins.step_wrapper",
     "Script": "task.script",
     "ScriptConfig": "task.script",
     "Config": "task.task",
     "Task": "task.task",
     "collate": "utils.data.collate",
     "collate_non_null": "utils.data.collate",
+    "get_named_leaves": "utils.debugging",
     "BaseFileDownloader": "utils.experiments",
+    "ContextTimer": "utils.experiments",
     "CumulativeTimer": "utils.experiments",
     "DataDownloader": "utils.experiments",
     "IntervalTicker": "utils.experiments",
@@ -198,6 +250,7 @@ NAME_MAP: dict[str, str] = {
     "stage_environment": "utils.experiments",
     "to_markdown_table": "utils.experiments",
     "jit": "utils.jax",
+    "save_jaxpr_dot": "utils.jaxpr",
     "ColoredFormatter": "utils.logging",
     "configure_logging": "utils.logging",
     "one_hot": "utils.numpy",
@@ -207,8 +260,13 @@ NAME_MAP: dict[str, str] = {
     "compute_nan_ratio": "utils.pytree",
     "flatten_array": "utils.pytree",
     "flatten_pytree": "utils.pytree",
+    "pytree_has_nans": "utils.pytree",
+    "reshuffle_pytree": "utils.pytree",
+    "reshuffle_pytree_along_dims": "utils.pytree",
+    "reshuffle_pytree_independently": "utils.pytree",
     "slice_array": "utils.pytree",
     "slice_pytree": "utils.pytree",
+    "update_pytree": "utils.pytree",
     "TextBlock": "utils.text",
     "camelcase_to_snakecase": "utils.text",
     "colored": "utils.text",
@@ -235,9 +293,12 @@ NAME_MAP.update(
         "LOG_ERROR_SUMMARY": "utils.logging",
         "LOG_PING": "utils.logging",
         "LOG_STATUS": "utils.logging",
+        "NormType": "nn.norm",
         "Output": "task.mixins.output",
         "Phase": "core.state",
         "RawConfigType": "task.base",
+        "ActivationFunction": "nn.equinox",
+        "DTYPE": "nn.equinox",
     },
 )
 
@@ -275,7 +336,27 @@ 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.equinox import (
+        DTYPE,
+        ActivationFunction,
+        MLPHyperParams,
+        export_eqx_mlp,
+        load_eqx,
+        load_eqx_mlp,
+        make_eqx_mlp,
+        save_eqx,
+    )
+    from xax.nn.export import (
+        export,
+        export_flax,
+        export_with_params,
+    )
+    from xax.nn.geom import (
+        euler_to_quat,
+        get_projected_gravity_vector_from_quat,
+        quat_to_euler,
+    )
+    from xax.nn.norm import NormType, cast_norm_type, get_norm
     from xax.nn.parallel import is_master
     from xax.task.base import RawConfigType
     from xax.task.launchers.base import BaseLauncher
@@ -290,12 +371,15 @@ if IMPORT_ALL or TYPE_CHECKING:
     from xax.task.mixins.cpu_stats import CPUStatsOptions
     from xax.task.mixins.data_loader import DataloaderConfig
     from xax.task.mixins.gpu_stats import GPUStatsOptions
+    from xax.task.mixins.step_wrapper import StepContext
     from xax.task.mixins.train import Batch, Output
     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 get_named_leaves
     from xax.utils.experiments import (
         BaseFileDownloader,
+        ContextTimer,
         CumulativeTimer,
         DataDownloader,
         IntervalTicker,
@@ -319,6 +403,7 @@ if IMPORT_ALL or TYPE_CHECKING:
         to_markdown_table,
     )
     from xax.utils.jax import jit
+    from xax.utils.jaxpr import save_jaxpr_dot
     from xax.utils.logging import (
         LOG_ERROR_SUMMARY,
         LOG_PING,
@@ -332,8 +417,13 @@ if IMPORT_ALL or TYPE_CHECKING:
         compute_nan_ratio,
         flatten_array,
         flatten_pytree,
+        pytree_has_nans,
+        reshuffle_pytree,
+        reshuffle_pytree_along_dims,
+        reshuffle_pytree_independently,
         slice_array,
         slice_pytree,
+        update_pytree,
     )
     from xax.utils.text import (
         TextBlock,

xax/nn/equinox.py

@@ -0,0 +1,180 @@
+"""Equinox utilities."""
+
+import json
+import logging
+from pathlib import Path
+from typing import Callable, Literal, TypedDict, cast
+
+import equinox as eqx
+import jax
+from jaxtyping import PRNGKeyArray
+
+logger = logging.getLogger(__name__)
+
+ActivationFunction = Literal[
+    "relu",
+    "tanh",
+    "celu",
+    "elu",
+    "gelu",
+    "glu",
+    "hard_sigmoid",
+    "hard_silu",
+    "hard_swish",
+    "hard_tanh",
+    "leaky_relu",
+    "log_sigmoid",
+    "log_softmax",
+    "logsumexp",
+    "relu6",
+    "selu",
+    "sigmoid",
+    "soft_sign",
+    "softmax",
+    "softplus",
+    "sparse_plus",
+    "sparse_sigmoid",
+    "silu",
+    "swish",
+    "squareplus",
+    "mish",
+    "identity",
+]
+
+DTYPE = Literal["float32", "float64"]
+
+DTYPE_MAP: dict[DTYPE, jax.numpy.dtype] = {
+    "float32": jax.numpy.float32,
+    "float64": jax.numpy.float64,
+}
+
+
+class MLPHyperParams(TypedDict):
+    """Hyperparameters of an Equinox MLP."""
+
+    in_size: int | Literal["scalar"]
+    out_size: int | Literal["scalar"]
+    width_size: int
+    depth: int
+    activation: ActivationFunction
+    final_activation: ActivationFunction
+    use_bias: bool
+    use_final_bias: bool
+    dtype: DTYPE
+
+
+def _infer_activation(activation: ActivationFunction) -> Callable:
+    if activation == "identity":
+        return lambda x: x
+    try:
+        return getattr(jax.nn, activation)
+    except AttributeError:
+        raise ValueError(f"Activation function `{activation}` not found in `jax.nn`")
+
+
+def make_eqx_mlp(hyperparams: MLPHyperParams, key: PRNGKeyArray = jax.random.PRNGKey(0)) -> eqx.nn.MLP:
+    """Create an Equinox MLP from a set of hyperparameters.
+
+    Args:
+        hyperparams: The hyperparameters of the MLP.
+        key: The PRNG key to use for the MLP.
+    """
+    activation = _infer_activation(hyperparams["activation"])
+    final_activation = _infer_activation(hyperparams["final_activation"])
+    dtype = DTYPE_MAP[hyperparams["dtype"]]
+
+    return eqx.nn.MLP(
+        in_size=hyperparams["in_size"],
+        out_size=hyperparams["out_size"],
+        width_size=hyperparams["width_size"],
+        depth=hyperparams["depth"],
+        activation=activation,
+        final_activation=final_activation,
+        use_bias=hyperparams["use_bias"],
+        use_final_bias=hyperparams["use_final_bias"],
+        dtype=dtype,
+        key=key,
+    )
+
+
+def export_eqx_mlp(
+    model: eqx.nn.MLP,
+    output_path: str | Path,
+    dtype: jax.numpy.dtype = eqx._misc.default_floating_dtype(),
+) -> None:
+    """Serialize an Equinox MLP to a .eqx file.
+
+    Args:
+        model: The JAX MLP to export.
+        output_path: The path to save the exported model.
+        dtype: The dtype of the model.
+    """
+    activation = model.activation.__name__
+    final_activation = model.final_activation.__name__
+
+    if final_activation == "<lambda>":
+        logger.warning("Final activation is a lambda function. Assuming identity.")
+        final_activation = "identity"
+
+    # cast strings to ActivationFunction for type checking
+    activation = cast(ActivationFunction, activation)
+    final_activation = cast(ActivationFunction, final_activation)
+
+    if dtype not in DTYPE_MAP.values():
+        raise ValueError(f"Invalid dtype: {dtype}. Must be one of {DTYPE_MAP.values()}")
+
+    dtype = {v: k for k, v in DTYPE_MAP.items()}[dtype]
+
+    hyperparams: MLPHyperParams = {
+        "in_size": model.in_size,
+        "out_size": model.out_size,
+        "width_size": model.width_size,
+        "depth": model.depth,
+        "activation": activation,
+        "final_activation": final_activation,
+        "use_bias": model.use_bias,
+        "use_final_bias": model.use_final_bias,
+        "dtype": dtype,
+    }
+
+    with open(output_path, "wb") as f:
+        hyperparam_str = json.dumps(hyperparams)
+        f.write((hyperparam_str + "\n").encode(encoding="utf-8"))
+        eqx.tree_serialise_leaves(f, model)
+
+
+def save_eqx(
+    model: eqx.Module,
+    output_path: str | Path,
+) -> None:
+    """Serialize an Equinox module to a .eqx file.
+
+    Args:
+        model: The Equinox module to export.
+        output_path: The path to save the exported model.
+    """
+    with open(output_path, "wb") as f:
+        eqx.tree_serialise_leaves(f, model)
+
+
+def load_eqx(
+    model: eqx.Module,
+    eqx_file: str | Path,
+) -> eqx.Module:
+    """Deserialize an Equinox module from a .eqx file.
+
+    Args:
+        model: The Equinox module to load into.
+        eqx_file: The path to the .eqx file to load.
+    """
+    with open(eqx_file, "rb") as f:
+        return eqx.tree_deserialise_leaves(f, model)
+
+
+def load_eqx_mlp(
+    eqx_file: str | Path,
+) -> eqx.nn.MLP:
+    with open(eqx_file, "rb") as f:
+        hyperparams = json.loads(f.readline().decode(encoding="utf-8"))
+        model = make_eqx_mlp(hyperparams=hyperparams)
+        return eqx.tree_deserialise_leaves(f, model)

xax/nn/export.py

@@ -0,0 +1,147 @@
+"""Export JAX functions to TensorFlow SavedModel format."""
+
+import logging
+from pathlib import Path
+from typing import Callable
+
+import flax
+import jax
+import tensorflow as tf
+from jax.experimental import jax2tf
+from jaxtyping import Array, PyTree
+from orbax.export import ExportManager, JaxModule, ServingConfig
+
+logger = logging.getLogger(__name__)
+
+
+def _run_infer(tf_module: tf.Module, input_shapes: list[tuple[int, ...]], batch_size: int | None) -> tf.Tensor:
+    """Warm up the model by running it once."""
+    if batch_size is not None:
+        test_inputs = [
+            jax.random.normal(jax.random.PRNGKey(42), (batch_size, *input_shape)) for input_shape in input_shapes
+        ]
+    else:
+        test_inputs = [jax.random.normal(jax.random.PRNGKey(42), (1, *input_shape)) for input_shape in input_shapes]
+    if not hasattr(tf_module, "infer"):
+        raise ValueError("Model does not have an infer method")
+    return tf_module.infer(*test_inputs)
+
+
+def export(
+    model: Callable,
+    input_shapes: list[tuple[int, ...]],
+    output_dir: str | Path = "export",
+    batch_size: int | None = None,
+) -> None:
+    """Export a JAX function to TensorFlow SavedModel.
+
+    Note: Tensorflow GraphDef can't be larger than 2GB - https://github.com/tensorflow/tensorflow/issues/51870
+    You can avoid this by saving model parameters as non-constants.
+
+    Args:
+        model: The JAX function to export.
+        input_shapes: The shape of the input tensors, excluding batch dimension.
+        output_dir: Directory to save the exported model.
+        batch_size: Optional batch dimension. If None, a polymorphic batch dimension is used.
+    """
+    tf_module = tf.Module()
+    # Create a polymorphic shape specification for each input
+    poly_spec = "(b, ...)" if batch_size is not None else "(None, ...)"
+    polymorphic_shapes = [poly_spec] * len(input_shapes)
+    tf_module.infer = tf.function(  # type: ignore [attr-defined]
+        jax2tf.convert(
+            model,
+            polymorphic_shapes=polymorphic_shapes,
+            # setting this to False will allow the model to run on platforms other than the one that exports the model
+            # https://github.com/jax-ml/jax/blob/051687dc4c899df3d95c30b812ade401d8b31166/jax/experimental/jax2tf/README.md?plain=1#L1342
+            # generally though I think native_serialization is recommended
+            native_serialization=False,
+            with_gradient=False,
+        ),
+        autograph=False,
+        input_signature=[tf.TensorSpec([batch_size] + list(input_shape), tf.float32) for input_shape in input_shapes],
+    )
+
+    # warm up the model
+    _run_infer(tf_module, input_shapes, batch_size)
+
+    logger.info("Exporting SavedModel to %s", output_dir)
+    tf.saved_model.save(
+        tf_module,
+        output_dir,
+    )
+
+
+def export_with_params(
+    model: Callable,
+    params: PyTree,
+    input_shapes: list[tuple[int, ...]],
+    output_dir: str | Path = "export",
+    batch_dim: int | None = None,
+) -> None:
+    """Export a JAX function that takes parameters to TensorFlow SavedModel.
+
+    Args:
+        model: The JAX function to export. Should take parameters as first argument.
+        params: The parameters to use for the model.
+        input_shapes: The shape of the input tensors, excluding batch dimension.
+        output_dir: Directory to save the exported model.
+        batch_dim: Optional batch dimension. If None, a polymorphic batch dimension is used.
+    """
+    param_vars = tf.nest.map_structure(tf.Variable, params)
+
+    converted_model = jax2tf.convert(model)
+
+    def model_fn(*inputs: PyTree) -> Array:
+        return converted_model(param_vars, *inputs)
+
+    tf_module = tf.Module()
+    tf_module._variables = tf.nest.flatten(param_vars)  # type: ignore [attr-defined]
+    tf_module.infer = tf.function(  # type: ignore [attr-defined]
+        model_fn,
+        jit_compile=True,
+        autograph=False,
+        input_signature=[tf.TensorSpec([batch_dim] + list(input_shape), tf.float32) for input_shape in input_shapes],
+    )
+
+    # warm up the model
+    _run_infer(tf_module, input_shapes, batch_dim)
+
+    logger.info("Exporting SavedModel to %s", output_dir)
+    tf.saved_model.save(tf_module, output_dir)
+
+
+def export_flax(
+    model: flax.linen.Module,
+    params: PyTree,
+    input_shape: tuple[int, ...],
+    preprocessor: Callable | None = None,
+    postprocessor: Callable | None = None,
+    input_name: str = "inputs",
+    output_name: str = "outputs",
+    output_dir: str | Path = "export",
+) -> None:
+    jax_module = JaxModule(
+        params, model.apply, trainable=False, input_polymorphic_shape="(b, ...)"
+    )  # if you want to use a batch dimension
+
+    # to avoid mapping sequences to ambiguous mappings
+    if postprocessor is None:
+
+        def postprocessor(x: PyTree) -> PyTree:
+            return {output_name: x}
+
+    export_manager = ExportManager(
+        jax_module,
+        [
+            ServingConfig(
+                "serving_default",
+                input_signature=[tf.TensorSpec([None] + list(input_shape), tf.float32, name=input_name)],
+                tf_preprocessor=preprocessor,
+                tf_postprocessor=postprocessor,
+            )
+        ],
+    )
+
+    logger.info("Exporting model to %s", output_dir)
+    export_manager.save(output_dir)

xax/nn/geom.py

@@ -73,3 +73,29 @@ def euler_to_quat(euler_3: jax.Array) -> jax.Array:
     quat = quat / jnp.linalg.norm(quat, axis=-1, keepdims=True)
 
     return quat
+
+
+def get_projected_gravity_vector_from_quat(quat: jax.Array, eps: float = 1e-6) -> jax.Array:
+    """Calculates the gravity vector projected onto the local frame given a quaternion orientation.
+
+    Args:
+        quat: A quaternion (w,x,y,z) representing the orientation, shape (*, 4).
+        eps: A small epsilon value to avoid division by zero.
+
+    Returns:
+        A 3D vector representing the gravity in the local frame, shape (*, 3).
+    """
+    # Normalize quaternion
+    quat = quat / (jnp.linalg.norm(quat, axis=-1, keepdims=True) + eps)
+    w, x, y, z = jnp.split(quat, 4, axis=-1)
+
+    # Gravity vector in world frame is [0, 0, -1] (pointing down)
+    # Rotate gravity vector using quaternion rotation
+
+    # Calculate quaternion rotation: q * [0,0,-1] * q^-1
+    gx = 2 * (x * z - w * y)
+    gy = 2 * (y * z + w * x)
+    gz = w * w - x * x - y * y + z * z
+
+    # Note: We're rotating [0,0,-1], so we negate gz to match the expected direction
+    return jnp.concatenate([gx, gy, -gz], axis=-1)

xax/nn/norm.py

@@ -0,0 +1,23 @@
+"""Normalization utilities."""
+
+from typing import Literal, cast, get_args
+
+import jax.numpy as jnp
+
+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: jnp.ndarray, norm: NormType) -> jnp.ndarray:
+    match norm:
+        case "l1":
+            return jnp.abs(x)
+        case "l2":
+            return jnp.square(x)
+        case _:
+            raise ValueError(f"Invalid norm: {norm}")

xax/requirements.txt

@@ -1,6 +1,7 @@
 # requirements.txt
 
 # Core ML/JAX dependencies
+attrs
 jax
 jaxtyping
 equinox

xax/task/base.py

@@ -81,6 +81,12 @@ class BaseTask(Generic[Config]):
     def on_training_end(self, state: State) -> State:
         return state
 
+    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__