xax 0.2.17__tar.gz → 0.2.19__tar.gz

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

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

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

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

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

@@ -92,7 +92,11 @@ class BaseTask(Generic[Config]):
 
     @functools.cached_property
     def task_path(self) -> Path:
-        return Path(inspect.getfile(self.__class__))
+        try:
+            return Path(inspect.getfile(self.__class__))
+        except OSError:
+            logger.warning("Could not resolve task path for %s, returning current working directory")
+            return Path.cwd()
 
     @functools.cached_property
     def task_module(self) -> str:
@@ -172,14 +176,18 @@ class BaseTask(Generic[Config]):
         Returns:
             The merged configs.
         """
-        task_path = Path(inspect.getfile(cls))
+        try:
+            task_path = Path(inspect.getfile(cls))
+        except OSError:
+            logger.warning("Could not resolve task path for %s, returning current working directory", cls.__name__)
+            task_path = Path.cwd()
         cfg = OmegaConf.structured(cls.get_config_class())
         cfg = OmegaConf.merge(cfg, *(get_config(other_cfg, task_path) for other_cfg in cfgs))
         if use_cli:
             args = use_cli if isinstance(use_cli, list) else sys.argv[1:]
             if "-h" in args or "--help" in args:
-                sys.stderr.write(OmegaConf.to_yaml(cfg))
-                sys.stderr.flush()
+                sys.stdout.write(OmegaConf.to_yaml(cfg, sort_keys=True))
+                sys.stdout.flush()
                 sys.exit(0)
 
             # Attempts to load any paths as configs.

{xax-0.2.17 → xax-0.2.19}/xax/task/mixins/artifacts.py

@@ -43,8 +43,14 @@ class ArtifactsMixin(BaseTask[Config]):
     def run_dir(self) -> Path:
         run_dir = get_run_dir()
         if run_dir is None:
-            task_file = inspect.getfile(self.__class__)
-            run_dir = Path(task_file).resolve().parent
+            try:
+                task_file = inspect.getfile(self.__class__)
+                run_dir = Path(task_file).resolve().parent
+            except OSError:
+                logger.warning(
+                    "Could not resolve task path for %s, returning current working directory", self.__class__.__name__
+                )
+                run_dir = Path.cwd()
         return run_dir / self.task_name
 
     @property

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

@@ -6,7 +6,7 @@ import logging
 import tarfile
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Generic, Literal, Sequence, TypeVar, cast, overload
+from typing import Generic, Literal, Self, Sequence, TypeVar, cast, overload
 
 import equinox as eqx
 import jax
@@ -46,7 +46,6 @@ class CheckpointingConfig(ArtifactsConfig):
     save_every_n_seconds: float | None = field(60.0 * 60.0, help="Save a checkpoint every N seconds")
     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")
 
 
 Config = TypeVar("Config", bound=CheckpointingConfig)
@@ -306,3 +305,11 @@ class CheckpointingMixin(ArtifactsMixin[Config], Generic[Config]):
         self.on_after_checkpoint_save(ckpt_path, state)
 
         return ckpt_path
+
+    @classmethod
+    def load_config(cls, ckpt_path: str | Path) -> Config:
+        return cls.get_config(load_ckpt(Path(ckpt_path), part="config"), use_cli=False)
+
+    @classmethod
+    def load_task(cls, ckpt_path: str | Path) -> Self:
+        return cls(cls.load_config(ckpt_path))

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

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