xax 0.1.8__py3-none-any.whl → 0.1.9__py3-none-any.whl

xax/__init__.py

@@ -12,7 +12,7 @@ and running the update script:
     python -m scripts.update_api --inplace
 """
 
-__version__ = "0.1.8"
+__version__ = "0.1.9"
 
 # This list shouldn't be modified by hand; instead, run the update script.
 __all__ = [
@@ -23,7 +23,6 @@ __all__ = [
     "get_run_dir",
     "load_user_config",
     "State",
-    "cast_phase",
     "FourierEmbeddings",
     "IdentityPositionalEmbeddings",
     "LearnedPositionalEmbeddings",
@@ -41,9 +40,6 @@ __all__ = [
     "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",
@@ -180,7 +176,6 @@ NAME_MAP: dict[str, str] = {
     "get_run_dir": "core.conf",
     "load_user_config": "core.conf",
     "State": "core.state",
-    "cast_phase": "core.state",
     "FourierEmbeddings": "nn.embeddings",
     "IdentityPositionalEmbeddings": "nn.embeddings",
     "LearnedPositionalEmbeddings": "nn.embeddings",
@@ -198,9 +193,6 @@ NAME_MAP: dict[str, str] = {
     "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",
@@ -329,7 +321,7 @@ if IMPORT_ALL or TYPE_CHECKING:
         get_run_dir,
         load_user_config,
     )
-    from xax.core.state import Phase, State, cast_phase
+    from xax.core.state import Phase, State
     from xax.nn.embeddings import (
         EmbeddingKind,
         FourierEmbeddings,
@@ -354,11 +346,6 @@ if IMPORT_ALL or TYPE_CHECKING:
         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,

xax/core/state.py

@@ -1,9 +1,10 @@
 """Defines a dataclass for keeping track of the current training state."""
 
 import time
-from dataclasses import dataclass
-from typing import Literal, NotRequired, TypedDict, cast, get_args
+from dataclasses import asdict, dataclass
+from typing import Any, Literal, NotRequired, TypedDict, Unpack, cast
 
+import jax
 from omegaconf import MISSING
 
 from xax.core.conf import field
@@ -11,12 +12,6 @@ from xax.core.conf import field
 Phase = Literal["train", "valid"]
 
 
-def cast_phase(raw_phase: str) -> Phase:
-    args = get_args(Phase)
-    assert raw_phase in args, f"Invalid phase: '{raw_phase}' Valid options are {args}"
-    return cast(Phase, raw_phase)
-
-
 class StateDict(TypedDict, total=False):
     num_steps: NotRequired[int]
     num_samples: NotRequired[int]
@@ -24,10 +19,11 @@ class StateDict(TypedDict, total=False):
     num_valid_samples: NotRequired[int]
     start_time_s: NotRequired[float]
     elapsed_time_s: NotRequired[float]
-    raw_phase: NotRequired[str]
+    phase: NotRequired[Phase]
 
 
-@dataclass
+@jax.tree_util.register_dataclass
+@dataclass(frozen=True, kw_only=True)
 class State:
     num_steps: int = field(MISSING, help="Number of steps so far")
     num_samples: int = field(MISSING, help="Number of sample so far")
@@ -35,15 +31,11 @@ class State:
     num_valid_samples: int = field(MISSING, help="Number of validation samples so far")
     start_time_s: float = field(MISSING, help="Start time of training")
     elapsed_time_s: float = field(MISSING, help="Total elapsed time so far")
-    raw_phase: str = field(MISSING, help="Current training phase")
+    _phase: int = field(MISSING, help="Current training phase")
 
     @property
     def phase(self) -> Phase:
-        return cast_phase(self.raw_phase)
-
-    @phase.setter
-    def phase(self, phase: Phase) -> None:
-        self.raw_phase = phase
+        return cast(Phase, ["train", "valid"][self._phase])
 
     @classmethod
     def init_state(cls) -> "State":
@@ -54,7 +46,7 @@ class State:
             num_valid_samples=0,
             start_time_s=time.time(),
             elapsed_time_s=0.0,
-            raw_phase="train",
+            _phase=0,
         )
 
     @property
@@ -69,3 +61,16 @@ class State:
                 return self.num_valid_steps
             case _:
                 raise ValueError(f"Invalid phase: {phase}")
+
+    def replace(self, **kwargs: Unpack[StateDict]) -> "State":
+        extra_kwargs: dict[str, Any] = {}  # noqa: ANN401
+        if "phase" in kwargs:
+            phase = kwargs.pop("phase")
+            match phase:
+                case "train":
+                    extra_kwargs["_phase"] = 0
+                case "valid":
+                    extra_kwargs["_phase"] = 1
+                case _:
+                    raise ValueError(f"Invalid phase: {phase}")
+        return State(**{**asdict(self), **kwargs, **extra_kwargs})

xax/task/base.py

@@ -16,7 +16,8 @@ from types import TracebackType
 from typing import Generic, Self, TypeVar, cast
 
 import jax
-from omegaconf import Container, DictConfig, OmegaConf
+from omegaconf import DictConfig, OmegaConf
+from omegaconf.base import SCMode
 
 from xax.core.state import State
 from xax.utils.text import camelcase_to_snakecase
@@ -66,9 +67,6 @@ class BaseTask(Generic[Config]):
 
         self.config = config
 
-        if isinstance(self.config, Container):
-            OmegaConf.resolve(self.config)
-
     def on_step_start(self, state: State) -> State:
         return state
 
@@ -195,7 +193,15 @@ class BaseTask(Generic[Config]):
                 cfg = OmegaConf.merge(cfg, *(get_config(path, task_path) for path in paths))
             cfg = OmegaConf.merge(cfg, OmegaConf.from_cli(non_paths))
 
-        return cast(Config, cfg)
+        return cast(
+            Config,
+            OmegaConf.to_container(
+                cfg,
+                resolve=True,
+                throw_on_missing=True,
+                structured_config_mode=SCMode.INSTANTIATE,
+            ),
+        )
 
     @classmethod
     def config_str(cls, *cfgs: RawConfigType, use_cli: bool | list[str] = True) -> str:

xax/task/mixins/train.py

@@ -53,6 +53,7 @@ from xax.utils.experiments import (
     get_packages_with_versions,
     get_training_code,
 )
+from xax.utils.jax import jit as xax_jit
 from xax.utils.logging import LOG_STATUS
 from xax.utils.text import highlight_exception_message, show_info
 
@@ -212,30 +213,27 @@ class TrainMixin(
 
     def on_step_end(self, state: State) -> State:
         state = super().on_step_end(state)
-        state.elapsed_time_s = time.time() - state.start_time_s
-        return state
+        return state.replace(elapsed_time_s=time.time() - state.start_time_s)
 
-    def log_train_step(self, model: PyTree, batch: Batch, output: Output, state: State) -> None:
+    def log_train_step(self, batch: Batch, output: Output, state: State) -> None:
         """Override this function to do logging during the training phase.
 
         This function is called after the model forward pass and before the
         backward pass. It is called in the training phase.
 
         Args:
-            model: The current model.
             batch: The batch from the dataloader.
             output: The model output.
             state: The current training state.
         """
 
-    def log_valid_step(self, model: PyTree, batch: Batch, output: Output, state: State) -> None:
+    def log_valid_step(self, batch: Batch, output: Output, state: State) -> None:
         """Override this function to do logging during the validation phase.
 
         This function is called after the model forward pass. It is called in
         the validation phase.
 
         Args:
-            model: The current model.
             batch: The batch from the dataloader.
             output: The model output.
             state: The current training state.
@@ -248,7 +246,7 @@ class TrainMixin(
             for k, v in d.items():
                 self.logger.log_scalar(k, v, namespace=ns)
 
-    def log_step(self, model: PyTree, batch: Batch, output: Output, loss: Array, state: State) -> None:
+    def log_step(self, batch: Batch, output: Output, loss: Array, state: State) -> None:
         phase = state.phase
 
         self.logger.log_scalar("loss", loss, namespace="loss")
@@ -257,9 +255,9 @@ class TrainMixin(
         # Delegate to the appropriate logging function based on the phase.
         match phase:
             case "train":
-                self.log_train_step(model, batch, output, state)
+                self.log_train_step(batch, output, state)
             case "valid":
-                self.log_valid_step(model, batch, output, state)
+                self.log_valid_step(batch, output, state)
             case _:
                 raise KeyError(f"Unknown phase: {phase}")
 
@@ -332,8 +330,7 @@ class TrainMixin(
 
         return model, optimizer, opt_state, state
 
-    @eqx.filter_jit
-    def get_output(self, model: PyTree, batch: Batch) -> Output:
+    def get_output(self, model: PyTree, batch: Batch, state: State) -> Output:
         """Gets the output from the model.
 
         By default, we assume the model is a function that takes the batch as
@@ -343,11 +340,11 @@ class TrainMixin(
         Args:
             model: The current model.
             batch: The current minibatch of samples.
+            state: The current training state.
         """
         raise NotImplementedError("`get_output` must be implemented by the subclass")
 
-    @eqx.filter_jit
-    def compute_loss(self, model: PyTree, batch: Batch, output: Output) -> Array:
+    def compute_loss(self, model: PyTree, batch: Batch, output: Output, state: State) -> Array:
         """Gets the loss for the current batch.
 
         By default, we assume the model is a function that takes the batch as
@@ -358,6 +355,7 @@ class TrainMixin(
             model: The current model.
             batch: The current minibatch of samples.
             output: The output from the model.
+            state: The current training state.
 
         Returns:
             The computed loss, as a tensor.
@@ -366,24 +364,32 @@ class TrainMixin(
             raise ValueError(f"When model output is not the loss, you must override `compute_loss`. Got {type(output)}")
         return output
 
-    @eqx.filter_jit
-    def get_output_and_loss(self, model: PyTree, batch: Batch) -> tuple[Array, Output]:
-        output = self.get_output(model, batch)
-        loss = self.compute_loss(model, batch, output)
+    def get_output_and_loss(
+        self,
+        model_static: PyTree,
+        model_arr: PyTree,
+        batch: Batch,
+        state: State,
+    ) -> tuple[Array, Output]:
+        model = eqx.combine(model_arr, model_static)
+        output = self.get_output(model, batch, state)
+        loss = self.compute_loss(model, batch, output, state)
         return loss, output
 
-    @eqx.filter_jit
     def update(
         self,
-        model: PyTree,
+        model_static: PyTree,
+        model_arr: PyTree,
         optimizer: optax.GradientTransformation,
         opt_state: optax.OptState,
         batch: Batch,
+        state: State,
     ) -> tuple[Array, PyTree, optax.OptState, Output]:
-        (loss, output), grads = eqx.filter_value_and_grad(self.get_output_and_loss, has_aux=True)(model, batch)
-        updates, opt_state = optimizer.update(grads, opt_state, model)
-        model = eqx.apply_updates(model, updates)
-        return loss, model, opt_state, output
+        grad_fn = eqx.filter_value_and_grad(self.get_output_and_loss, has_aux=True)
+        (loss, output), grads = grad_fn(model_static, model_arr, batch, state)
+        updates, opt_state = optimizer.update(grads, opt_state, model_arr)
+        model_arr = eqx.apply_updates(model_arr, updates)
+        return loss, model_arr, opt_state, output
 
     def get_size_of_batch(self, batch: Batch) -> int | None:
         """Gets the batch size for the current batch.
@@ -457,21 +463,31 @@ class TrainMixin(
         self.logger.log_file("training_code.txt", get_training_code(self))
         self.logger.log_file("config.yaml", self.config_str(self.config, use_cli=False))
 
-    @eqx.filter_jit
+    def model_partition_fn(self, item: Any) -> bool:  # noqa: ANN401
+        return eqx.is_inexact_array(item)
+
+    @xax_jit(static_argnames=["self", "model_static", "optimizer"])
     def train_step(
         self,
-        model: PyTree,
+        model_static: PyTree,
+        model_arr: PyTree,
         optimizer: optax.GradientTransformation,
         opt_state: optax.OptState,
         batch: Batch,
+        state: State,
     ) -> tuple[PyTree, optax.OptState, Array, Output]:
-        loss, model, opt_state, output = self.update(model, optimizer, opt_state, batch)
-        return model, opt_state, loss, output
+        loss, model_arr, opt_state, output = self.update(model_static, model_arr, optimizer, opt_state, batch, state)
+        return model_arr, opt_state, loss, output
 
-    @eqx.filter_jit
-    def val_step(self, model: PyTree, batch: Batch) -> tuple[PyTree, Array, Output]:
-        loss, output = eqx.filter_jit(self.get_output_and_loss)(model, batch)
-        return model, loss, output
+    @xax_jit(static_argnames=["self", "model_static"])
+    def val_step(
+        self,
+        model_static: PyTree,
+        model_arr: PyTree,
+        batch: Batch,
+        state: State,
+    ) -> tuple[Array, Output]:
+        return self.get_output_and_loss(model_static, model_arr, batch, state)
 
     def train_loop(
         self,
@@ -482,36 +498,46 @@ class TrainMixin(
         valid_pf: Iterator[Batch],
         state: State,
     ) -> None:
+        model_arr, model_static = eqx.partition(model, self.model_partition_fn)
+
         while not self.is_training_over(state):
             if self.valid_step_timer.is_valid_step(state):
                 valid_batch = next(valid_pf)
-                with self.step_context("model_step"):
-                    model, loss, output = self.val_step(model, valid_batch)
+                state = state.replace(
+                    phase="valid",
+                    num_valid_steps=state.num_valid_steps + 1,
+                    num_valid_samples=state.num_valid_samples + (self.get_size_of_batch(valid_batch) or 0),
+                )
 
-                # Perform logging.
-                with self.step_context("write_logs"):
-                    state.phase = "valid"
-                    self.log_step(model, valid_batch, output, loss, state)
-                    state.num_valid_samples += 1
+                loss, output = self.val_step(model_static, model_arr, valid_batch, state)
+                self.log_step(valid_batch, output, loss, state)
 
             state = self.on_step_start(state)
-
-            with self.step_context("model_step"):
-                train_batch = next(train_pf)
-                model, opt_state, loss, output = self.train_step(model, optimizer, opt_state, train_batch)
-
-            with self.step_context("write_logs"):
-                state.phase = "train"
-                self.log_step(model, train_batch, output, loss, state)
-                state.num_steps += 1
-                state.num_samples += self.get_size_of_batch(train_batch) or 0
+            train_batch = next(train_pf)
+            state = state.replace(
+                phase="train",
+                num_steps=state.num_steps + 1,
+                num_samples=state.num_samples + (self.get_size_of_batch(train_batch) or 0),
+            )
+
+            model_arr, opt_state, loss, output = self.train_step(
+                model_static=model_static,
+                model_arr=model_arr,
+                optimizer=optimizer,
+                opt_state=opt_state,
+                batch=train_batch,
+                state=state,
+            )
+            self.log_step(train_batch, output, loss, state)
 
             state = self.on_step_end(state)
 
             if self.should_checkpoint(state):
+                model = eqx.combine(model_arr, model_static)
                 self.save_checkpoint(model, optimizer, opt_state, state)
 
         # After finishing training, save the final checkpoint.
+        model = eqx.combine(model_arr, model_static)
         self.save_checkpoint(model, optimizer, opt_state, state)
 
     @contextlib.contextmanager

xax/task/script.py

@@ -3,6 +3,8 @@
 from dataclasses import dataclass
 from typing import Generic, TypeVar
 
+import jax
+
 from xax.task.base import BaseConfig, BaseTask
 from xax.task.mixins import (
     ArtifactsConfig,
@@ -20,6 +22,7 @@ from xax.task.mixins import (
 )
 
 
+@jax.tree_util.register_dataclass
 @dataclass(kw_only=True)
 class ScriptConfig(
     CPUStatsConfig,

{xax-0.1.8.dist-info → xax-0.1.9.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: xax
-Version: 0.1.8
+Version: 0.1.9
 Summary: A library for fast Jax experimentation
 Home-page: https://github.com/kscalelabs/xax
 Author: Benjamin Bolte

{xax-0.1.8.dist-info → xax-0.1.9.dist-info}/RECORD

@@ -1,10 +1,10 @@
-xax/__init__.py,sha256=9WNjoeAF7enu7YXQqshpVG1FucGdSkxwrRa-ELDDuUs,13713
+xax/__init__.py,sha256=_xb60-jl7arZEleSwUw4ElPaq4MzD24_ZYQrnWO5_cs,13391
 xax/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/requirements-dev.txt,sha256=qkscNkFzWd1S5fump-AKH53rR65v2x5FmboFdy_kKvs,128
 xax/requirements.txt,sha256=9LAEZ5c5gqRSARRVA6xJsVTa4MebPZuC4yOkkwkZJFw,297
 xax/core/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/core/conf.py,sha256=Wuo5WLRWuRTgb8eaihvnG_NZskTu0-P3JkIcl_hKINM,5124
-xax/core/state.py,sha256=y123fL7pMgk25TPG6KN0LRIF_eYnD9eP7OfqtoQJGNE,2178
+xax/core/state.py,sha256=WwW0qDm-be9MMOT-bGWEFvaWF4iq2FP9xRSn1zq_4A8,2507
 xax/nn/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/nn/embeddings.py,sha256=bQGxBFxkLwi2MQLkRfGaHPH5P_KKB21HdI7VNWTKIOQ,11847
 xax/nn/equinox.py,sha256=5fdOKRXqAVZPsV-aEez3i1wamr_oBYnG74GP1jEthjM,4843
@@ -14,9 +14,9 @@ xax/nn/geom.py,sha256=eK7I8fUHBc3FT7zpm5Yf__bXFQ4LtX6sa17-DxojLTo,3202
 xax/nn/norm.py,sha256=cDmYf5CtyzmuCiWdSP5nr8nZKQOmaZueDQXMPnThg6c,548
 xax/nn/parallel.py,sha256=fnTiT7MsG7eQrJvqwjIz2Ifo3P27TuxIJzmpGYSa_dQ,4608
 xax/task/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-xax/task/base.py,sha256=4fUjrG-llQpeESQuaQbww4M6WR6djjTK89fY20UV9zU,7610
+xax/task/base.py,sha256=E4l1yCrAkM2TVTbVYrmk6BoVHMkbD4IYsTT921XOyi0,7760
 xax/task/logger.py,sha256=1SZjVC6UCtZUoMPcpp3ckotL324QDeYDvHVhf5MHVqg,36271
-xax/task/script.py,sha256=zt36Sobdoer86gXHqc4sMAW7bqZRVl6IEExuQZH2USk,926
+xax/task/script.py,sha256=bMMIJoUtpSBvPp6-7bejTrajTXvSg0794sYLKdPIToE,972
 xax/task/task.py,sha256=UHMpnv__gqMcfbC_L-Hhk-DCnUYlFVsgbNf-v8o8B7U,1424
 xax/task/launchers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/task/launchers/base.py,sha256=8LB_r6YISKu1vq1zk3aVYmiedRr9MxE5IMRocs6unFI,731
@@ -39,7 +39,7 @@ xax/task/mixins/logger.py,sha256=6oXsJJyNUx6YT3q58FVXMZBUpMgjVkGre6BXFN20cVI,280
 xax/task/mixins/process.py,sha256=d1opVgvc6bOFXb7R58b07F4P5lbSZIzYaajtE0eBbpw,1477
 xax/task/mixins/runnable.py,sha256=IYIsLd2k09g-_y6o44EhJqT7E6BpsyEMmsyLSuzqjtc,1979
 xax/task/mixins/step_wrapper.py,sha256=-Yu5Nft2CRw1JvZt6J_94SM1vqX8fk08IDK95Pmd2ew,1648
-xax/task/mixins/train.py,sha256=vsH_QpyrThlh9AzWnyvDJv58Y8U_516oi8gmMq_0iMg,22333
+xax/task/mixins/train.py,sha256=JbrSiBqpgOrdDanNYuAzzh2radPrXOVrHYA6VcxjIzY,23248
 xax/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/utils/debugging.py,sha256=9WlCrEqbq-SVXPEM4rhsLYERH97XNX7XSYLSI3sgKGk,1619
 xax/utils/experiments.py,sha256=5CUja1H_cx4dnVqTGQekOpIhqISwHtAgLxZ34GV7cwM,29229
@@ -56,8 +56,8 @@ 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=ZCMGfSfr2_b2qZbq9ywPD0zej5tpVSId2JftXpwfB5k,4686
 xax/utils/types/hashable_array.py,sha256=l5iIcFmkYzfGeaZmcSoeFkthFASqM8xJYK3AXhZQYwc,992
-xax-0.1.8.dist-info/licenses/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
-xax-0.1.8.dist-info/METADATA,sha256=wnBSNRByXJzgQPuZqNWooidfFdqcT4w8gbwlBgzbJk8,1877
-xax-0.1.8.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-xax-0.1.8.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
-xax-0.1.8.dist-info/RECORD,,
+xax-0.1.9.dist-info/licenses/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
+xax-0.1.9.dist-info/METADATA,sha256=Ou8KmYWWNxgo_9ZAU2KLaeGeXAxd6b9qJ95ky4HRm-o,1877
+xax-0.1.9.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
+xax-0.1.9.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
+xax-0.1.9.dist-info/RECORD,,