xax 0.2.18__py3-none-any.whl → 0.2.19__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.18"
+__version__ = "0.2.19"
 
 # This list shouldn't be modified by hand; instead, run the update script.
 __all__ = [

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

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

{xax-0.2.18.dist-info → xax-0.2.19.dist-info}/RECORD

@@ -1,4 +1,4 @@
-xax/__init__.py,sha256=ru-amUXLIpFI1cE7NfkPpff6NnuLbSh6LHbJkmvH2zM,15733
+xax/__init__.py,sha256=Pv4UWs4GGvojcnBDz4hk9whnU4ZEklWXmTMaTwlSZPM,15733
 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
@@ -16,7 +16,7 @@ xax/nn/metrics.py,sha256=OAkeScwhi-wTBIJ59KHUhYbZTq4V4V-LG-mKlxMJ7bY,3238
 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
-xax/task/base.py,sha256=ymGld72zTSw50PBpX4WwqQDyJgQSagDrLh56aqvLT40,7720
+xax/task/base.py,sha256=qF-ed7bFrBIRHxM7hVPLTGRYYwiLp98H0b6NWF80LTQ,8101
 xax/task/logger.py,sha256=W_BpluYvQai1lh1dDCAj-2_mWUC1buhwJncHygDffjc,41125
 xax/task/script.py,sha256=bMMIJoUtpSBvPp6-7bejTrajTXvSg0794sYLKdPIToE,972
 xax/task/task.py,sha256=UHMpnv__gqMcfbC_L-Hhk-DCnUYlFVsgbNf-v8o8B7U,1424
@@ -31,7 +31,7 @@ xax/task/loggers/state.py,sha256=0Jy0NYnY4c0qt0LvNlaTaCKOSqk5SCKln5VdyuQGnIc,140
 xax/task/loggers/stdout.py,sha256=giKSW2R83YkgRefm3BLkE7t8Pbj5Dux4AgsdJxYIbGo,6619
 xax/task/loggers/tensorboard.py,sha256=sRyBbeBeVXDTYhPZIKIapW0JEfL9hqqzhNTeIcSd374,8883
 xax/task/mixins/__init__.py,sha256=D3oU31rB9FeOr9MPLleLt5JFbftUr4sBTwgnwQdc2qA,809
-xax/task/mixins/artifacts.py,sha256=Ma7fwsp-SA1w6GcuBSskszj5TB83yxYJm4Ns_EnqkI4,3018
+xax/task/mixins/artifacts.py,sha256=IBPAQMCGd7PQiZHfSjLakPW5j7cNuL6AsW6QkVSc02E,3277
 xax/task/mixins/checkpointing.py,sha256=w-xsx8tH8PPYRf9eoqvz-4kAdunHMvazcoZAIj6niyI,11468
 xax/task/mixins/compile.py,sha256=PG5aF3W9v_xGiImHgUJ7gmwuQQoSQWufdpl2N_mlLX0,3922
 xax/task/mixins/cpu_stats.py,sha256=rO_9a82ZdsNec61ya4FpYE-rWqPhpijRSXsOfc6caFA,9595
@@ -58,8 +58,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=ebtHENhyUzSjyJTlbMaLtcckQIJ7EtgJiok_40TJZpo,4689
 xax/utils/types/hashable_array.py,sha256=l5iIcFmkYzfGeaZmcSoeFkthFASqM8xJYK3AXhZQYwc,992
-xax-0.2.18.dist-info/licenses/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
-xax-0.2.18.dist-info/METADATA,sha256=phU79eE74vMM79HouJP9iQgTDONJ59no3ktUl5GHF78,1880
-xax-0.2.18.dist-info/WHEEL,sha256=wXxTzcEDnjrTwFYjLPcsW_7_XihufBwmpiBeiXNBGEA,91
-xax-0.2.18.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
-xax-0.2.18.dist-info/RECORD,,
+xax-0.2.19.dist-info/licenses/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
+xax-0.2.19.dist-info/METADATA,sha256=fmZSs-aNP8Rwo9bPOnb9Fiu9UUlZ_5mxXnVl4fZkDMw,1880
+xax-0.2.19.dist-info/WHEEL,sha256=wXxTzcEDnjrTwFYjLPcsW_7_XihufBwmpiBeiXNBGEA,91
+xax-0.2.19.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
+xax-0.2.19.dist-info/RECORD,,