haliax 1.4.dev292__tar.gz → 1.4.dev294__tar.gz

{haliax-1.4.dev292 → haliax-1.4.dev294}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: haliax
-Version: 1.4.dev292
+Version: 1.4.dev294
 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/

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

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

{haliax-1.4.dev292 → haliax-1.4.dev294}/src/haliax/partitioning.py

@@ -3,6 +3,7 @@ import functools
 import threading
 import typing
 import warnings
+from itertools import chain
 from math import prod
 from typing import Callable, ContextManager, Mapping, Optional, ParamSpec, Sequence, TypeVar, Union
 
@@ -38,6 +39,7 @@ class ResourceAxis(StringHolderEnum):
 
     MODEL = "model"
     DATA = "data"
+    REPLICA = "replica"
 
 
 class _ResourceMappingHolder:
@@ -584,7 +586,17 @@ def sharding_for_axis(
 def pspec_for_axis(axis: AxisSelection, mapping: Optional[ResourceMapping] = None) -> PartitionSpec:
     """Get the PartitionSpec for a single axis"""
     axis = ensure_tuple(axis)
-    return PartitionSpec(*(physical_axis_name(a, mapping) for a in axis))
+    phys_axes = []
+    for a in axis:
+        pa = physical_axis_name(a, mapping)
+        if pa is None or isinstance(pa, str):
+            phys_axes.append(pa)
+        else:
+            # I have no way to resolve the mypy check :)
+            for i in pa:
+                phys_axes.append(i)
+
+    return PartitionSpec(*phys_axes)
 
 
 def round_axis_for_partitioning(axis: Axis, mapping: Optional[ResourceMapping] = None) -> Axis:

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

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