xax 0.2.20__py3-none-any.whl → 0.2.21__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.20"
+__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/cli/edit_config.py

@@ -52,14 +52,24 @@ def main() -> None:
             print(colored(line, "light-cyan"), flush=True)
 
     # Saves the edited config to the checkpoint.
-    with tarfile.open(args.ckpt_path, "w:gz") as tar:
+    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)
 
-        def add_file_bytes(name: str, data: bytes) -> None:  # noqa: ANN401
-            info = tarfile.TarInfo(name=name)
-            info.size = len(data)
-            tar.addfile(info, io.BytesIO(data))
+        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_file_bytes("config", edited_config_str.encode())
+            # 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__":

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.20.dist-info → xax-0.2.21.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: xax
-Version: 0.2.20
+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.20.dist-info → xax-0.2.21.dist-info}/RECORD

@@ -1,20 +1,18 @@
-xax/__init__.py,sha256=yWQcHMlP2cKIpfJCJLXv796F-AAHQhS-1sRxu7871mw,15733
+xax/__init__.py,sha256=R1U4pGfRfz5IfEp6iJjjxliG8oDtnzldJ0BKnm36CBE,15207
 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
 xax/cli/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-xax/cli/edit_config.py,sha256=99x_k6aNimbcebi2vSJhln-cv4364h6GQdRccuv_qcs,2069
+xax/cli/edit_config.py,sha256=LQUIlOS6hvPZyVEaMme3FP-62M0BKQPYavCwVDWuBLw,2600
 xax/core/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/core/conf.py,sha256=d7Dp_GwKnaxtkztlSrJSM_LR0UYJX_FWTtceIWCBkxc,5138
 xax/core/state.py,sha256=KsNMnM_RgsZ2Ntc2pp4Fi6zG4rZb_89-kqmyGxDvyRg,4974
 xax/nn/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/nn/embeddings.py,sha256=bQGxBFxkLwi2MQLkRfGaHPH5P_KKB21HdI7VNWTKIOQ,11847
-xax/nn/equinox.py,sha256=JZuSApD4bL0UK5W1nrQtucWYvNWUha07J6LTLk_RX-Y,4910
-xax/nn/export.py,sha256=pRfM2B4hB2EvljysC6AjtgB_7Cn7JtaP3dhYU2stZtY,5545
 xax/nn/functions.py,sha256=bA5kJYzMtFM8eUqBC086i355zJMAO7k_vPFNSDBI9-s,2814
 xax/nn/geom.py,sha256=A7WPefMvgwUNReZC7_HX1GmvHPASyghbaXaKsuhwDrE,7382
 xax/nn/losses.py,sha256=Q_NVnm5n4UPBvp5nI_1aUptfXnqFYoUeFwySiyvopHg,272
-xax/nn/metrics.py,sha256=OAkeScwhi-wTBIJ59KHUhYbZTq4V4V-LG-mKlxMJ7bY,3238
+xax/nn/metrics.py,sha256=zuvPXlRQczBTLHD4ilNGmZaiq6Yie3rxCMq6JkI_kos,3154
 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
@@ -60,9 +58,9 @@ 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.20.dist-info/licenses/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
-xax-0.2.20.dist-info/METADATA,sha256=YCcDox7HsHIVUeDeZYCvoaGYhsU7TzspX-v-Xw0-H4g,1880
-xax-0.2.20.dist-info/WHEEL,sha256=0CuiUZ_p9E4cD6NyLD6UG80LBXYyiSYZOKDm5lp32xk,91
-xax-0.2.20.dist-info/entry_points.txt,sha256=uRC6rx5ce0bf-FblJaZSBMxxKFfMyoWTf8OWbBmLSe8,61
-xax-0.2.20.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
-xax-0.2.20.dist-info/RECORD,,
+xax-0.2.21.dist-info/licenses/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
+xax-0.2.21.dist-info/METADATA,sha256=cPa4omiKYq0HgWjZ-NddSZ432cKfrLaHBz3CrA1l9pk,1247
+xax-0.2.21.dist-info/WHEEL,sha256=DnLRTWE75wApRYVsjgc6wsVswC54sMSJhAEd4xhDpBk,91
+xax-0.2.21.dist-info/entry_points.txt,sha256=uRC6rx5ce0bf-FblJaZSBMxxKFfMyoWTf8OWbBmLSe8,61
+xax-0.2.21.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
+xax-0.2.21.dist-info/RECORD,,

{xax-0.2.20.dist-info → xax-0.2.21.dist-info}/WHEEL

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

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/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)