xax 0.1.2__tar.gz → 0.1.3__tar.gz

{xax-0.1.2/xax.egg-info → xax-0.1.3}/PKG-INFO

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

{xax-0.1.2 → xax-0.1.3}/xax/__init__.py

@@ -12,7 +12,7 @@ and running the update script:
     python -m scripts.update_api --inplace
 """
 
-__version__ = "0.1.2"
+__version__ = "0.1.3"
 
 # This list shouldn't be modified by hand; instead, run the update script.
 __all__ = [

{xax-0.1.2 → xax-0.1.3}/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-0.1.2 → xax-0.1.3}/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-0.1.2 → xax-0.1.3}/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
@@ -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.2 → xax-0.1.3/xax.egg-info}/PKG-INFO

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