haliax 1.4.dev326__tar.gz → 1.4.dev327__tar.gz

{haliax-1.4.dev326 → haliax-1.4.dev327}/PKG-INFO

@@ -1,11 +1,12 @@
-Metadata-Version: 2.3
+Metadata-Version: 2.4
 Name: haliax
-Version: 1.4.dev326
+Version: 1.4.dev327
 Summary: Named Tensors for Legible Deep Learning in JAX
 Project-URL: Homepage, https://github.com/stanford-crfm/haliax
 Project-URL: Bug Tracker, https://github.com/stanford-crfm/haliax/issues/
 Project-URL: Documentation, https://haliax.readthedocs.io/en/latest/
 Author-email: David Hall <dlwh@cs.stanford.edu>
+License-File: LICENSE
 Classifier: Development Status :: 4 - Beta
 Classifier: Intended Audience :: Science/Research
 Classifier: License :: OSI Approved :: Apache Software License

haliax-1.4.dev327/src/haliax/__about__.py

@@ -0,0 +1 @@
+__version__ = "1.4.dev327"

{haliax-1.4.dev326 → haliax-1.4.dev327}/src/haliax/partitioning.py

@@ -8,7 +8,7 @@ from typing import Callable, ContextManager, Mapping, Optional, ParamSpec, Seque
 
 import equinox as eqx
 import jax
-from equinox import module_update_wrapper
+from equinox import is_array, module_update_wrapper
 from jax.lax import with_sharding_constraint
 from jax.sharding import Mesh, NamedSharding, PartitionSpec, SingleDeviceSharding
 from jaxtyping import PyTree
@@ -20,7 +20,7 @@ from .axis import Axis, AxisSelection, AxisSelector
 from .core import NamedArray
 from .jax_utils import Static, is_in_jit, is_jax_array_like, is_on_mac_metal
 from .tree_util import hashable_combine, hashable_partition
-from .util import StringHolderEnum, ensure_tuple, is_named_array
+from .util import StringHolderEnum, ensure_tuple
 
 
 PhysicalAxisSpec = Union[(str), Sequence[str]]
@@ -274,7 +274,7 @@ class _NamedJitWrapper(eqx.Module):
         if out_axis_resources is None:
             out_axis_resources = axis_resources
 
-        dynamic_argspec, static_argspec = hashable_partition((args, kwargs), is_jax_array_like)
+        dynamic_argspec, static_argspec = hashable_partition((args, kwargs), is_array)
         dynamic = (self._dynamic_fun, dynamic_argspec)
 
         donate_args = self._donate_args
@@ -436,7 +436,7 @@ def named_jit(
             **pjit_args,
         )
 
-    dynamic_fun, static_fun = hashable_partition(fn, is_jax_array_like)
+    dynamic_fun, static_fun = hashable_partition(fn, is_array)
 
     wrapper = _NamedJitWrapper(
         fn,
@@ -514,7 +514,7 @@ def _named_pjit_cache(fun_names, **jitkwargs) -> WrappedCallable:
         fun = hashable_combine(dynamic_fun, static_fun)
         args, kwargs = hashable_combine(dynamic_spec, static_spec)
         out = fun(*args, **kwargs)
-        out_dynamic, out_static = hashable_partition(out, is_jax_array_like)
+        out_dynamic, out_static = hashable_partition(out, is_array)
         return out_dynamic, Static(out_static)
 
     fun_name, fun_qualname = fun_names
@@ -543,7 +543,7 @@ def _cached_filter_eval_shape(fun, *args, **kwargs):
     eval_shape is surprisingly expensive, so we cache it. We use this for named_pjit for evaluating resource partitions
     of the output.
     """
-    dynamic, static = hashable_partition((fun, args, kwargs), is_jax_array_like)
+    dynamic, static = hashable_partition((fun, args, kwargs), is_array)
     if static not in _eval_shape_cache:
         _eval_shape_cache[static] = eqx.filter_eval_shape(fun, *args, **kwargs)
 

{haliax-1.4.dev326 → haliax-1.4.dev327}/src/haliax/quantization.py

@@ -10,7 +10,6 @@ from typing import Optional, Protocol, TypeVar
 import equinox as eqx
 import jax
 from jax import numpy as jnp
-from jax._src.tree_util import BuiltInKeyEntry
 from jax.tree_util import DictKey, FlattenedIndexKey, GetAttrKey, SequenceKey
 from jax.typing import DTypeLike
 
@@ -253,7 +252,7 @@ def _matches_target_fp8(key_path, config: Fp8Config) -> bool:
     return re.match(config.targets, key_path_str) is not None
 
 
-def _key_path_to_str(key_path: tuple[BuiltInKeyEntry, ...]) -> str:
+def _key_path_to_str(key_path: tuple) -> str:
     out = ""
     for k in key_path:
         match k:

haliax-1.4.dev326/src/haliax/__about__.py

@@ -1 +0,0 @@
-__version__ = "1.4.dev326"