haliax 1.4.dev295__tar.gz → 1.4.dev296__tar.gz

{haliax-1.4.dev295 → haliax-1.4.dev296}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: haliax
-Version: 1.4.dev295
+Version: 1.4.dev296
 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.dev296/src/haliax/__about__.py

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

{haliax-1.4.dev295 → haliax-1.4.dev296}/src/haliax/partitioning.py

@@ -3,7 +3,6 @@ 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
 
@@ -586,17 +585,7 @@ 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)
-    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)
+    return PartitionSpec(*(physical_axis_name(a, mapping) for a in axis))
 
 
 def round_axis_for_partitioning(axis: Axis, mapping: Optional[ResourceMapping] = None) -> Axis:

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

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