xax 0.4.0__tar.gz → 0.4.2__tar.gz

{xax-0.4.0/xax.egg-info → xax-0.4.2}/PKG-INFO

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

{xax-0.4.0 → xax-0.4.2}/xax/__init__.py

@@ -12,7 +12,7 @@ and running the update script:
     python -m scripts.update_api --inplace
 """
 
-__version__ = "0.4.0"
+__version__ = "0.4.2"
 
 # This list shouldn't be modified by hand; instead, run the update script.
 __all__ = [
@@ -171,6 +171,7 @@ __all__ = [
     "uncolored",
     "wrapped",
     "FrozenDict",
+    "freeze_dict",
     "HashableArray",
     "hashable_array",
 ]
@@ -371,6 +372,7 @@ NAME_MAP: dict[str, str] = {
     "uncolored": "utils.text",
     "wrapped": "utils.text",
     "FrozenDict": "utils.types.frozen_dict",
+    "freeze_dict": "utils.types.frozen_dict",
     "HashableArray": "utils.types.hashable_array",
     "hashable_array": "utils.types.hashable_array",
 }
@@ -572,7 +574,7 @@ if IMPORT_ALL or TYPE_CHECKING:
         uncolored,
         wrapped,
     )
-    from xax.utils.types.frozen_dict import FrozenDict
+    from xax.utils.types.frozen_dict import FrozenDict, freeze_dict
     from xax.utils.types.hashable_array import HashableArray, hashable_array
 
 del TYPE_CHECKING, IMPORT_ALL

{xax-0.4.0 → xax-0.4.2}/xax/utils/experiments.py

@@ -5,6 +5,7 @@ import datetime
 import enum
 import functools
 import hashlib
+import importlib.metadata
 import inspect
 import itertools
 import json
@@ -29,7 +30,6 @@ from typing import Any, Iterator, Mapping, Self, Sequence, 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
@@ -464,7 +464,7 @@ def get_packages_with_versions() -> str:
     Returns:
         A dictionary of packages and their versions.
     """
-    packages = [(pkg.key, pkg.version) for pkg in pkg_resources.working_set]
+    packages = [(dist.metadata["Name"], dist.version) for dist in importlib.metadata.distributions()]
     return "\n".join([f"{key}=={version}" for key, version in sorted(packages)])
 
 

{xax-0.4.0 → xax-0.4.2}/xax/utils/types/frozen_dict.py

@@ -146,3 +146,7 @@ def unfreeze(x: FrozenDict[K, V] | dict[str, Any]) -> dict[Any, Any]:  # noqa: A
         return ys
     else:
         return x
+
+
+def freeze_dict(x: Mapping[K, V]) -> FrozenDict[K, V]:
+    return FrozenDict(x)

{xax-0.4.0 → xax-0.4.2/xax.egg-info}/PKG-INFO

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