xax 0.2.19__tar.gz → 0.2.21__tar.gz

{xax-0.2.19/xax.egg-info → xax-0.2.21}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: xax
-Version: 0.2.19
+Version: 0.2.21
 Summary: A library for fast Jax experimentation
 Home-page: https://github.com/kscalelabs/xax
 Author: Benjamin Bolte
@@ -31,22 +31,6 @@ 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: exportable
-Requires-Dist: flax; extra == "exportable"
-Requires-Dist: orbax-export; extra == "exportable"
-Requires-Dist: tensorflow; extra == "exportable"
-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: flax; extra == "all"
-Requires-Dist: orbax-export; extra == "all"
-Requires-Dist: tensorflow; extra == "all"
 Dynamic: author
 Dynamic: description
 Dynamic: description-content-type

{xax-0.2.19 → xax-0.2.21}/pyproject.toml

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

{xax-0.2.19 → xax-0.2.21}/setup.py

@@ -14,12 +14,6 @@ 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] = [
-    "flax",
-    "orbax-export",
-    "tensorflow",
-]
-
 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"
@@ -39,8 +33,6 @@ setup(
     tests_require=requirements_dev,
     extras_require={
         "dev": requirements_dev,
-        "exportable": requirements_export,
-        "all": requirements_dev + requirements_export,
     },
     package_data={
         "xax": [
@@ -48,4 +40,9 @@ setup(
             "requirements*.txt",
         ],
     },
+    entry_points={
+        "console_scripts": [
+            "xax-edit-config=xax.cli.edit_config:main",
+        ],
+    },
 )

{xax-0.2.19 → xax-0.2.21}/xax/__init__.py

@@ -12,7 +12,7 @@ and running the update script:
     python -m scripts.update_api --inplace
 """
 
-__version__ = "0.2.19"
+__version__ = "0.2.21"
 
 # This list shouldn't be modified by hand; instead, run the update script.
 __all__ = [
@@ -34,12 +34,6 @@ __all__ = [
     "get_positional_embeddings",
     "get_rotary_embeddings",
     "rotary_embeddings",
-    "MLPHyperParams",
-    "export_eqx_mlp",
-    "load_eqx",
-    "load_eqx_mlp",
-    "make_eqx_mlp",
-    "save_eqx",
     "cubic_bezier_interpolation",
     "euler_to_quat",
     "get_projected_gravity_vector_from_quat",
@@ -215,12 +209,6 @@ 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",
     "cubic_bezier_interpolation": "nn.geom",
     "euler_to_quat": "nn.geom",
     "get_projected_gravity_vector_from_quat": "nn.geom",
@@ -392,16 +380,6 @@ if IMPORT_ALL or TYPE_CHECKING:
         get_rotary_embeddings,
         rotary_embeddings,
     )
-    from xax.nn.equinox import (
-        DTYPE,
-        ActivationFunction,
-        MLPHyperParams,
-        export_eqx_mlp,
-        load_eqx,
-        load_eqx_mlp,
-        make_eqx_mlp,
-        save_eqx,
-    )
     from xax.nn.geom import (
         cubic_bezier_interpolation,
         euler_to_quat,

xax-0.2.21/xax/cli/edit_config.py

@@ -0,0 +1,77 @@
+"""Lets you edit a checkpoint config programmatically."""
+
+import argparse
+import difflib
+import io
+import os
+import subprocess
+import tarfile
+import tempfile
+from pathlib import Path
+
+from omegaconf import OmegaConf
+
+from xax.task.mixins.checkpointing import load_ckpt
+from xax.utils.text import colored, show_info
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("ckpt_path", type=Path)
+    args = parser.parse_args()
+
+    # Loads the config from the checkpoint.
+    config = load_ckpt(args.ckpt_path, part="config")
+    config_str = OmegaConf.to_yaml(config)
+
+    # Opens the user's preferred editor to edit the config.
+    with tempfile.NamedTemporaryFile(suffix=".yaml", delete=False) as f:
+        f.write(config_str.encode("utf-8"))
+        f.flush()
+        subprocess.run([os.environ.get("EDITOR", "vim"), f.name], check=True)
+
+    # Loads the edited config.
+    try:
+        edited_config = OmegaConf.load(f.name)
+        edited_config_str = OmegaConf.to_yaml(edited_config, sort_keys=True)
+    finally:
+        os.remove(f.name)
+
+    if edited_config_str == config_str:
+        show_info("No changes were made to the config.")
+        return
+
+    # Diffs the original and edited configs.
+    diff = difflib.ndiff(config_str.splitlines(), edited_config_str.splitlines())
+    for line in diff:
+        if line.startswith("+ "):
+            print(colored(line, "light-green"), flush=True)
+        elif line.startswith("- "):
+            print(colored(line, "light-red"), flush=True)
+        elif line.startswith("? "):
+            print(colored(line, "light-cyan"), flush=True)
+
+    # Saves the edited config to the checkpoint.
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        with tarfile.open(args.ckpt_path, "r:gz") as src_tar:
+            for member in src_tar.getmembers():
+                if member.name != "config":  # Skip the old config file
+                    src_tar.extract(member, tmp_dir)
+
+        with tarfile.open(args.ckpt_path, "w:gz") as tar:
+            for root, _, files in os.walk(tmp_dir):
+                for file in files:
+                    file_path = os.path.join(root, file)
+                    arcname = os.path.relpath(file_path, tmp_dir)
+                    tar.add(file_path, arcname=arcname)
+
+            # Add the new config file
+            info = tarfile.TarInfo(name="config")
+            config_bytes = edited_config_str.encode()
+            info.size = len(config_bytes)
+            tar.addfile(info, io.BytesIO(config_bytes))
+
+
+if __name__ == "__main__":
+    # python -m xax.cli.edit_config
+    main()

{xax-0.2.19 → xax-0.2.21}/xax/nn/metrics.py

@@ -7,8 +7,6 @@ import jax
 import jax.numpy as jnp
 from jaxtyping import Array
 
-from xax.utils.jax import jit as xax_jit
-
 NormType = Literal["l1", "l2"]
 
 
@@ -36,7 +34,6 @@ def dynamic_time_warping(distance_matrix_nm: Array) -> Array: ...
 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.
 

{xax-0.2.19 → xax-0.2.21}/xax/task/base.py

@@ -210,7 +210,7 @@ class BaseTask(Generic[Config]):
 
     @classmethod
     def config_str(cls, *cfgs: RawConfigType, use_cli: bool | list[str] = True) -> str:
-        return OmegaConf.to_yaml(cls.get_config(*cfgs, use_cli=use_cli))
+        return OmegaConf.to_yaml(cls.get_config(*cfgs, use_cli=use_cli), sort_keys=True)
 
     @classmethod
     def get_task(cls, *cfgs: RawConfigType, use_cli: bool | list[str] = True) -> Self:

{xax-0.2.19 → xax-0.2.21}/xax/task/mixins/checkpointing.py

@@ -292,7 +292,7 @@ class CheckpointingMixin(ArtifactsMixin[Config], Generic[Config]):
 
             if state is not None:
                 add_file_bytes("state", json.dumps(state.to_dict(), indent=2).encode())
-            add_file_bytes("config", OmegaConf.to_yaml(self.config).encode())
+            add_file_bytes("config", OmegaConf.to_yaml(self.config, sort_keys=True).encode())
 
         # Updates the symlink to the new checkpoint
         last_ckpt_path.unlink(missing_ok=True)

{xax-0.2.19 → xax-0.2.21/xax.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: xax
-Version: 0.2.19
+Version: 0.2.21
 Summary: A library for fast Jax experimentation
 Home-page: https://github.com/kscalelabs/xax
 Author: Benjamin Bolte
@@ -31,22 +31,6 @@ 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: exportable
-Requires-Dist: flax; extra == "exportable"
-Requires-Dist: orbax-export; extra == "exportable"
-Requires-Dist: tensorflow; extra == "exportable"
-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: flax; extra == "all"
-Requires-Dist: orbax-export; extra == "all"
-Requires-Dist: tensorflow; extra == "all"
 Dynamic: author
 Dynamic: description
 Dynamic: description-content-type

{xax-0.2.19 → xax-0.2.21}/xax.egg-info/SOURCES.txt

@@ -11,15 +11,16 @@ xax/requirements.txt
 xax.egg-info/PKG-INFO
 xax.egg-info/SOURCES.txt
 xax.egg-info/dependency_links.txt
+xax.egg-info/entry_points.txt
 xax.egg-info/requires.txt
 xax.egg-info/top_level.txt
+xax/cli/__init__.py
+xax/cli/edit_config.py
 xax/core/__init__.py
 xax/core/conf.py
 xax/core/state.py
 xax/nn/__init__.py
 xax/nn/embeddings.py
-xax/nn/equinox.py
-xax/nn/export.py
 xax/nn/functions.py
 xax/nn/geom.py
 xax/nn/losses.py

xax-0.2.21/xax.egg-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+xax-edit-config = xax.cli.edit_config:main

{xax-0.2.19 → xax-0.2.21}/xax.egg-info/requires.txt

@@ -14,19 +14,6 @@ tensorboard
 psutil
 requests
 
-[all]
-black
-darglint
-mypy
-ruff
-pytest
-types-pillow
-types-psutil
-types-requests
-flax
-orbax-export
-tensorflow
-
 [dev]
 black
 darglint
@@ -36,8 +23,3 @@ pytest
 types-pillow
 types-psutil
 types-requests
-
-[exportable]
-flax
-orbax-export
-tensorflow

xax-0.2.19/xax/nn/equinox.py

@@ -1,183 +0,0 @@
-"""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 as err:
-        raise ValueError(f"Activation function `{activation}` not found in `jax.nn`") from err
-
-
-def make_eqx_mlp(hyperparams: MLPHyperParams, *, key: PRNGKeyArray) -> 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 | None = None,
-) -> 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.
-    """
-    if dtype is None:
-        dtype = eqx._misc.default_floating_dtype()
-
-    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, key=jax.random.PRNGKey(0))
-        return eqx.tree_deserialise_leaves(f, model)

xax-0.2.19/xax/nn/export.py

@@ -1,154 +0,0 @@
-"""Export JAX functions to TensorFlow SavedModel format."""
-
-import logging
-from pathlib import Path
-from typing import Callable
-
-import jax
-from jaxtyping import Array, PyTree
-
-try:
-    import flax
-    import tensorflow as tf
-    from jax.experimental import jax2tf
-    from orbax.export import ExportManager, JaxModule, ServingConfig
-except ImportError as e:
-    raise ImportError(
-        "In order to export models, please install Xax with exportable dependencies, "
-        "using 'xax[exportable]` to install the required dependencies."
-    ) from e
-
-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)