xax 0.1.1__py3-none-any.whl → 0.1.3__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.1"
+__version__ = "0.1.3"
 
 # This list shouldn't be modified by hand; instead, run the update script.
 __all__ = [
@@ -66,6 +66,7 @@ __all__ = [
     "DataloaderConfig",
     "GPUStatsOptions",
     "StepContext",
+    "ValidStepTimer",
     "Script",
     "ScriptConfig",
     "Config",
@@ -219,6 +220,7 @@ NAME_MAP: dict[str, str] = {
     "DataloaderConfig": "task.mixins.data_loader",
     "GPUStatsOptions": "task.mixins.gpu_stats",
     "StepContext": "task.mixins.step_wrapper",
+    "ValidStepTimer": "task.mixins.train",
     "Script": "task.script",
     "ScriptConfig": "task.script",
     "Config": "task.task",
@@ -372,7 +374,7 @@ if IMPORT_ALL or TYPE_CHECKING:
     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.mixins.train import Batch, Output, ValidStepTimer
     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

xax/task/mixins/checkpointing.py

@@ -47,7 +47,6 @@ class CheckpointingConfig(ArtifactsConfig):
     only_save_most_recent: bool = field(True, help="Only keep the most recent checkpoint")
     load_from_ckpt_path: str | None = field(None, help="If set, load initial model weights from this path")
     load_ckpt_strict: bool = field(True, help="If set, only load weights for which have a matching key in the model")
-    save_tf_model: bool = field(False, help="If set, saves a Tensorflow version of the model")
 
 
 Config = TypeVar("Config", bound=CheckpointingConfig)
@@ -213,15 +212,6 @@ class CheckpointingMixin(ArtifactsMixin[Config], Generic[Config]):
             add_file("state", lambda buf: buf.write(json.dumps(asdict(state), indent=2).encode()))
             add_file("config", lambda buf: buf.write(OmegaConf.to_yaml(self.config).encode()))
 
-        if self.config.save_tf_model:
-            try:
-                from jax.experimental import jax2tf
-            except ModuleNotFoundError:
-                raise ImportError("Tensorflow is not installed. Install it with `pip install tensorflow`")
-
-            tf_model = jax2tf.convert(model)
-            add_file("model.tf", lambda buf: cloudpickle.dump(tf_model, buf))
-
         # Updates the symlink to the new checkpoint.
         last_ckpt_path.unlink(missing_ok=True)
         try:

xax/task/mixins/train.py

@@ -50,6 +50,7 @@ from xax.utils.experiments import (
     diff_configs,
     get_diff_string,
     get_git_state,
+    get_packages_with_versions,
     get_training_code,
 )
 from xax.utils.logging import LOG_STATUS
@@ -452,6 +453,7 @@ class TrainMixin(
         logger.log(LOG_STATUS, self.task_name)
         logger.log(LOG_STATUS, "JAX devices: %s", jax.devices())
         self.logger.log_file("git_state.txt", get_git_state(self))
+        self.logger.log_file("packages.txt", get_packages_with_versions())
         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))
 

xax/utils/experiments.py

@@ -28,6 +28,7 @@ from typing import Any, Iterator, Self, TypeVar, cast
 from urllib.parse import urlparse
 
 import git
+import pkg_resources
 import requests
 from jaxtyping import Array
 from omegaconf import MISSING, DictConfig, ListConfig, OmegaConf
@@ -109,8 +110,8 @@ class StateTimer:
 
     def step(self, state: State) -> None:
         cur_time = time.time()
-        self.step_timer.step(state.num_steps, cur_time)
-        self.sample_timer.step(state.num_samples, cur_time)
+        self.step_timer.step(state.num_steps if state.phase == "train" else state.num_valid_steps, cur_time)
+        self.sample_timer.step(state.num_samples if state.phase == "train" else state.num_valid_samples, cur_time)
         self.iter_timer.step(cur_time)
 
     def log_dict(self) -> dict[str, dict[str, int | float]]:
@@ -468,6 +469,16 @@ def get_git_state(obj: object) -> str:
         return traceback.format_exc()
 
 
+def get_packages_with_versions() -> str:
+    """Gets the packages and their versions.
+
+    Returns:
+        A dictionary of packages and their versions.
+    """
+    packages = [(pkg.key, pkg.version) for pkg in pkg_resources.working_set]
+    return "\n".join([f"{key}=={version}" for key, version in sorted(packages)])
+
+
 def get_training_code(obj: object) -> str:
     """Gets the text from the file containing the provided object.
 

{xax-0.1.1.dist-info → xax-0.1.3.dist-info}/METADATA

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

{xax-0.1.1.dist-info → xax-0.1.3.dist-info}/RECORD

@@ -1,4 +1,4 @@
-xax/__init__.py,sha256=JyKRACir9b0bkuG93bwxADFrVr-Lo76kenDBJtvb_wQ,13280
+xax/__init__.py,sha256=n7EXl0pwEPzlw2DjS-3ePgx0VoQnMDnHLVc5exkHGcM,13361
 xax/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/requirements-dev.txt,sha256=qkscNkFzWd1S5fump-AKH53rR65v2x5FmboFdy_kKvs,128
 xax/requirements.txt,sha256=9LAEZ5c5gqRSARRVA6xJsVTa4MebPZuC4yOkkwkZJFw,297
@@ -30,7 +30,7 @@ xax/task/loggers/stdout.py,sha256=bR0k-PfmFgLfPxLPb4hZw_8G_msA32UeHfAAu11nEYs,67
 xax/task/loggers/tensorboard.py,sha256=kI8LvBuBBhPgkP8TeaTQb9SQ0FqaIodwQh2SuWDCnIA,7706
 xax/task/mixins/__init__.py,sha256=D3oU31rB9FeOr9MPLleLt5JFbftUr4sBTwgnwQdc2qA,809
 xax/task/mixins/artifacts.py,sha256=2ezmZGzPGe3nhsd9KRkeHWWXdbT9m7drzimIfw6v1XY,2892
-xax/task/mixins/checkpointing.py,sha256=sRkVxJbQfqDf1-lp1KFrAGYWHhTlV8_DORxGQ_69P1A,8954
+xax/task/mixins/checkpointing.py,sha256=a6tVyISsDIz68rrhb1rAh3rjQlqkDVJCmSBmETQrnRM,8480
 xax/task/mixins/compile.py,sha256=FRsxwLnZjjxpeWJ7Bx_d8XUY50oDoGidgpeRt4ejeQk,3377
 xax/task/mixins/cpu_stats.py,sha256=C_t71UTrv4LwQzhO5iubsfomj4jYa9bzpE4zBcHdoHM,9211
 xax/task/mixins/data_loader.py,sha256=WjMWk9uACfBMMClLMcLPkE0WNIvlCZnmqyyqLqJpjX0,6545
@@ -39,10 +39,10 @@ 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=BEC7HSwBlGZDe7jCsedqEA8-K1Zx52-bTjsBONYIE5g,22225
+xax/task/mixins/train.py,sha256=8AaBXaopnrxtSZXldyFCE3QX1k5r3IsZMr6O0ICnNnU,22332
 xax/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/utils/debugging.py,sha256=9WlCrEqbq-SVXPEM4rhsLYERH97XNX7XSYLSI3sgKGk,1619
-xax/utils/experiments.py,sha256=_cwoBaiBxoQ_Tstm0rz7TEqfELqcktmPflb6AP1K0qA,28779
+xax/utils/experiments.py,sha256=5CUja1H_cx4dnVqTGQekOpIhqISwHtAgLxZ34GV7cwM,29229
 xax/utils/jax.py,sha256=tC0NNelbrSTzwNGluiwLGKtoHhVpgdzrv-xherB3VtY,4752
 xax/utils/jaxpr.py,sha256=S80nyEkv188RInzq3kCAdkQCU-bf6s0oPTrCE_LjkRs,2298
 xax/utils/logging.py,sha256=GAhTne2rdB4Fa1lzk06DMO15U8MTejn6XTClShC-ZtU,6622
@@ -53,8 +53,8 @@ xax/utils/tensorboard.py,sha256=_S70dS69pduiD05viHAGgYGsaBry1QL2ej6ZwUIXPOE,1617
 xax/utils/text.py,sha256=zo1sAoZe59GkpcpaHBVOQ0OekSMGXvOAyNa3lOJozCY,10628
 xax/utils/data/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/utils/data/collate.py,sha256=Rd9vMomr_S_zCa_Hi4dO-8ntzAfVwndIUtuXFA3iNcc,7066
-xax-0.1.1.dist-info/licenses/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
-xax-0.1.1.dist-info/METADATA,sha256=tJ4ilL3uBbykHBQTHbh-bN6m4hrHqivyyFeuI33ddX4,1877
-xax-0.1.1.dist-info/WHEEL,sha256=1tXe9gY0PYatrMPMDd6jXqjfpz_B-Wqm32CPfRC58XU,91
-xax-0.1.1.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
-xax-0.1.1.dist-info/RECORD,,
+xax-0.1.3.dist-info/licenses/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
+xax-0.1.3.dist-info/METADATA,sha256=m3AyjlRD9C-O2Tp5zH5i5TbEL7bZooeIpypUCYuYPtQ,1877
+xax-0.1.3.dist-info/WHEEL,sha256=DK49LOLCYiurdXXOXwGJm6U4DkHkg4lcxjhqwRa0CP4,91
+xax-0.1.3.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
+xax-0.1.3.dist-info/RECORD,,

{xax-0.1.1.dist-info → xax-0.1.3.dist-info}/WHEEL

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