xax 0.3.11__tar.gz → 0.3.12__tar.gz

{xax-0.3.11/xax.egg-info → xax-0.3.12}/PKG-INFO

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

{xax-0.3.11 → xax-0.3.12}/xax/__init__.py

@@ -12,7 +12,7 @@ and running the update script:
     python -m scripts.update_api --inplace
 """
 
-__version__ = "0.3.11"
+__version__ = "0.3.12"
 
 # This list shouldn't be modified by hand; instead, run the update script.
 __all__ = [
@@ -138,6 +138,7 @@ __all__ = [
     "worker_chunk",
     "profile",
     "compute_nan_ratio",
+    "diff_pytree",
     "flatten_array",
     "flatten_pytree",
     "get_pytree_mapping",
@@ -330,6 +331,7 @@ NAME_MAP: dict[str, str] = {
     "worker_chunk": "utils.numpy",
     "profile": "utils.profile",
     "compute_nan_ratio": "utils.pytree",
+    "diff_pytree": "utils.pytree",
     "flatten_array": "utils.pytree",
     "flatten_pytree": "utils.pytree",
     "get_pytree_mapping": "utils.pytree",
@@ -413,7 +415,12 @@ if IMPORT_ALL or TYPE_CHECKING:
         TransformerCache,
         TransformerStack,
     )
-    from xax.nn.distributions import Categorical, Distribution, MixtureOfGaussians, Normal
+    from xax.nn.distributions import (
+        Categorical,
+        Distribution,
+        MixtureOfGaussians,
+        Normal,
+    )
     from xax.nn.embeddings import (
         EmbeddingKind,
         FourierEmbeddings,
@@ -518,6 +525,7 @@ if IMPORT_ALL or TYPE_CHECKING:
     from xax.utils.profile import profile
     from xax.utils.pytree import (
         compute_nan_ratio,
+        diff_pytree,
         flatten_array,
         flatten_pytree,
         get_pytree_mapping,

{xax-0.3.11 → xax-0.3.12}/xax/nn/distributions.py

@@ -12,7 +12,6 @@ __all__ = [
     "MixtureOfGaussians",
 ]
 
-import math
 from abc import ABC, abstractmethod
 
 import jax
@@ -20,7 +19,7 @@ import jax.numpy as jnp
 from jaxtyping import Array, PRNGKeyArray
 
 STD_CLIP = 1e-6
-LOGIT_CLIP = math.log(1e4)
+LOGIT_CLIP = 6.0
 
 
 class Distribution(ABC):

{xax-0.3.11 → xax-0.3.12}/xax/utils/pytree.py

@@ -1,12 +1,15 @@
 """Utils for accessing, modifying, and otherwise manipulating pytrees."""
 
+from dataclasses import fields, is_dataclass
 from typing import Mapping, Sequence, TypeVar
 
 import chex
 import equinox as eqx
 import jax
 import jax.numpy as jnp
+import numpy as np
 from jax import Array
+from jax.core import get_aval
 from jaxtyping import PRNGKeyArray, PyTree
 
 T = TypeVar("T")
@@ -258,18 +261,79 @@ def tuple_insert(t: tuple[T, ...], index: int, value: T) -> tuple[T, ...]:
 def get_pytree_mapping(pytree: PyTree) -> dict[str, Array]:
     leaves: dict[str, Array] = {}
 
-    def _get_str(thing: PyTree) -> str:
-        if isinstance(thing, str):
-            return thing
-        if isinstance(thing, Sequence):
-            return "/".join(_get_str(x) for x in thing)
-        if isinstance(thing, Mapping):
-            return "/".join(f"{_get_str(k)}:{_get_str(v)}" for k, v in thing.items())
-        return str(thing)
-
     def _get_leaf(path: tuple, x: PyTree) -> None:
         if isinstance(x, jnp.ndarray):
-            leaves[_get_str(path)] = x
+            leaves[jax.tree_util.keystr(path, simple=True, separator="/")] = x
 
     jax.tree.map_with_path(_get_leaf, pytree)
     return leaves
+
+
+def diff_pytree(tree_a: PyTree, tree_b: PyTree, prefix: str = "") -> list[str]:
+    diffs = []
+
+    # Handles dataclasses.
+    if is_dataclass(tree_a) and is_dataclass(tree_b):
+        for field in fields(tree_a):
+            attr_a, attr_b = getattr(tree_a, field.name), getattr(tree_b, field.name)
+            diffs.extend(diff_pytree(attr_a, attr_b, prefix + f"{field.name}."))
+        return diffs
+
+    # Handle dict-like objects
+    elif isinstance(tree_a, Mapping) and isinstance(tree_b, Mapping):
+        if type(tree_a) is not type(tree_b):
+            diffs.append(f"{prefix}: type {type(tree_a)} vs {type(tree_b)}")
+            return diffs
+        keys_a, keys_b = set(tree_a.keys()), set(tree_b.keys())
+        for k in keys_a - keys_b:
+            diffs.append(f"{prefix}{k}: present in A only")
+        for k in keys_b - keys_a:
+            diffs.append(f"{prefix}{k}: present in B only")
+        for k in keys_a & keys_b:
+            diffs.extend(diff_pytree(tree_a[k], tree_b[k], prefix + f"{k}."))
+        return diffs
+
+    # Handle tuple/list
+    elif isinstance(tree_a, Sequence) and isinstance(tree_b, Sequence):
+        if type(tree_a) is not type(tree_b):
+            diffs.append(f"{prefix}: type {type(tree_a)} vs {type(tree_b)}")
+            return diffs
+        if len(tree_a) != len(tree_b):
+            diffs.append(f"{prefix}: different lengths {len(tree_a)} vs {len(tree_b)}")
+        for i, (a_i, b_i) in enumerate(zip(tree_a, tree_b, strict=True)):
+            diffs.extend(diff_pytree(a_i, b_i, prefix + f"[{i}]."))
+        return diffs
+
+    # Handles basic types.
+    elif isinstance(tree_a, (int, float, bool, str, type(None), np.number, np.bool, bytes)):
+        if tree_a != tree_b:
+            diffs.append(f"{prefix}: {tree_a!r} vs {tree_b!r}")
+        return diffs
+
+    # Handles Numpy arrays.
+    elif isinstance(tree_a, np.ndarray) and isinstance(tree_b, np.ndarray):
+        if tree_a.shape != tree_b.shape:
+            diffs.append(f"{prefix}: shape {tree_a.shape} vs {tree_b.shape}")
+        if tree_a.dtype != tree_b.dtype:
+            diffs.append(f"{prefix}: dtype {tree_a.dtype} vs {tree_b.dtype}")
+        return diffs
+
+    # Handle arrays (check shape/dtype)
+    elif isinstance(tree_a, jnp.ndarray) and isinstance(tree_b, jnp.ndarray):
+        if tree_a.shape != tree_b.shape:
+            diffs.append(f"{prefix}: shape {tree_a.shape} vs {tree_b.shape}")
+        if tree_a.dtype != tree_b.dtype:
+            diffs.append(f"{prefix}: dtype {tree_a.dtype} vs {tree_b.dtype}")
+        aval_a = get_aval(tree_a)
+        aval_b = get_aval(tree_b)
+        if aval_a != aval_b:  # pyright: ignore[reportAttributeAccessIssue]
+            diffs.append(f"{prefix}: aval {aval_a} vs {aval_b}")
+        return diffs
+
+    # Handle mismatched types
+    elif type(tree_a) is not type(tree_b):
+        diffs.append(f"{prefix}: type {type(tree_a)} vs {type(tree_b)}")
+        return diffs
+
+    else:
+        raise ValueError(f"Unknown type: {type(tree_a)}")

{xax-0.3.11 → xax-0.3.12/xax.egg-info}/PKG-INFO

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