xax 0.0.7__tar.gz → 0.1.0__tar.gz

{xax-0.0.7/xax.egg-info → xax-0.1.0}/PKG-INFO

@@ -1,12 +1,13 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.4
 Name: xax
-Version: 0.0.7
-Summary: The xax project
-Home-page: https://github.com/dpshai/xax
+Version: 0.1.0
+Summary: A library for fast Jax experimentation
+Home-page: https://github.com/kscalelabs/xax
 Author: Benjamin Bolte
 Requires-Python: >=3.11
 Description-Content-Type: text/markdown
 License-File: LICENSE
+Requires-Dist: attrs
 Requires-Dist: jax
 Requires-Dist: jaxtyping
 Requires-Dist: equinox
@@ -30,10 +31,28 @@ Requires-Dist: pytest; extra == "dev"
 Requires-Dist: types-pillow; extra == "dev"
 Requires-Dist: types-psutil; extra == "dev"
 Requires-Dist: types-requests; extra == "dev"
+Provides-Extra: export
+Requires-Dist: orbax-export; extra == "export"
+Requires-Dist: tensorflow; extra == "export"
+Provides-Extra: flax
+Requires-Dist: flax; extra == "flax"
+Provides-Extra: all
+Requires-Dist: black; extra == "all"
+Requires-Dist: darglint; extra == "all"
+Requires-Dist: mypy; extra == "all"
+Requires-Dist: ruff; extra == "all"
+Requires-Dist: pytest; extra == "all"
+Requires-Dist: types-pillow; extra == "all"
+Requires-Dist: types-psutil; extra == "all"
+Requires-Dist: types-requests; extra == "all"
+Requires-Dist: orbax-export; extra == "all"
+Requires-Dist: tensorflow; extra == "all"
+Requires-Dist: flax; extra == "all"
 Dynamic: author
 Dynamic: description
 Dynamic: description-content-type
 Dynamic: home-page
+Dynamic: license-file
 Dynamic: provides-extra
 Dynamic: requires-dist
 Dynamic: requires-python

{xax-0.0.7 → xax-0.1.0}/pyproject.toml

@@ -40,6 +40,7 @@ module = [
     "setuptools.*",
     "tensorboard.*",
     "transformers.*",
+    "orbax.export.*",
 ]
 
 ignore_missing_imports = true

{xax-0.0.7 → xax-0.1.0}/setup.py

@@ -14,6 +14,15 @@ with open("xax/requirements.txt", "r", encoding="utf-8") as f:
 with open("xax/requirements-dev.txt", "r", encoding="utf-8") as f:
     requirements_dev: list[str] = f.read().splitlines()
 
+requirements_export: list[str] = [
+    "orbax-export",
+    "tensorflow",
+]
+
+requirements_flax: list[str] = [
+    "flax",
+]
+
 with open("xax/__init__.py", "r", encoding="utf-8") as fh:
     version_re = re.search(r"^__version__ = \"([^\"]*)\"", fh.read(), re.MULTILINE)
 assert version_re is not None, "Could not find version in xax/__init__.py"
@@ -23,9 +32,9 @@ version: str = version_re.group(1)
 setup(
     name="xax",
     version=version,
-    description="The xax project",
+    description="A library for fast Jax experimentation",
     author="Benjamin Bolte",
-    url="https://github.com/dpshai/xax",
+    url="https://github.com/kscalelabs/xax",
     long_description=long_description,
     long_description_content_type="text/markdown",
     python_requires=">=3.11",
@@ -33,6 +42,9 @@ setup(
     tests_require=requirements_dev,
     extras_require={
         "dev": requirements_dev,
+        "export": requirements_export,
+        "flax": requirements_flax,
+        "all": requirements_dev + requirements_export + requirements_flax,
     },
     package_data={
         "xax": [

{xax-0.0.7 → xax-0.1.0}/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-0.1.0/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-0.1.0/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)