haliax 1.4.dev293__tar.gz → 1.4.dev295__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: haliax
-Version: 1.4.dev293
+Version: 1.4.dev295
 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/
@@ -67,8 +67,6 @@ please see the [Haliax tutorial](https://colab.research.google.com/drive/1TiTcQQ
 (We use the excellent [Equinox](https://github.com/patrick-kidger/equinox) library for its module system and tree transformations.)
 
 ```python
-import haliax.nn.normalization
-import haliax.nn.activations
 import equinox as eqx
 import jax
 import jax.numpy as jnp
@@ -90,7 +88,7 @@ def attention_scores(Key, KPos, query, key, mask):
         scores -= 1E9 * (1.0 - mask)
 
     # convert to probabilities
-    scores = haliax.nn.normalization.softmax(scores, KPos)
+    scores = haliax.nn.softmax(scores, KPos)
     return scores
 
 

{haliax-1.4.dev293 → haliax-1.4.dev295}/README.md

@@ -35,8 +35,6 @@ please see the [Haliax tutorial](https://colab.research.google.com/drive/1TiTcQQ
 (We use the excellent [Equinox](https://github.com/patrick-kidger/equinox) library for its module system and tree transformations.)
 
 ```python
-import haliax.nn.normalization
-import haliax.nn.activations
 import equinox as eqx
 import jax
 import jax.numpy as jnp
@@ -58,7 +56,7 @@ def attention_scores(Key, KPos, query, key, mask):
         scores -= 1E9 * (1.0 - mask)
 
     # convert to probabilities
-    scores = haliax.nn.normalization.softmax(scores, KPos)
+    scores = haliax.nn.softmax(scores, KPos)
     return scores
 
 

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

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

{haliax-1.4.dev293 → haliax-1.4.dev295}/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
 
@@ -585,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.dev293/src/haliax/__about__.py

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