haliax 1.4.dev294__tar.gz → 1.4.dev296__tar.gz

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

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

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

{haliax-1.4.dev294 → 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.dev294/src/haliax/__about__.py

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