haliax 1.4.dev381__tar.gz → 1.4.dev386__tar.gz

{haliax-1.4.dev381 → haliax-1.4.dev386}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: haliax
-Version: 1.4.dev381
+Version: 1.4.dev386
 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/
@@ -33,14 +33,14 @@ Description-Content-Type: text/markdown
 <a href="">
 <img alt="License" src="https://img.shields.io/github/license/stanford-crfm/haliax?color=blue" />
 </a>
-<a href="https://https://pypi.org/project/haliax/">
+<a href="https://pypi.org/project/haliax/">
     <img alt="PyPI" src="https://img.shields.io/pypi/v/haliax?color=blue" />
 </a>
 
 > *Though you don’t seem to be much for listening, it’s best to be careful. If you managed to catch hold of even just a piece of my name, you’d have all manner of power over me.*<br/>
 > — Patrick Rothfuss, *The Name of the Wind*
 
-Haliax is a [JAX](https:://github.com/google/jax) library for building neural networks with named tensors, in the tradition of Alexander Rush's [Tensor Considered Harmful](https://nlp.seas.harvard.edu/NamedTensor).
+Haliax is a [JAX](https://github.com/google/jax) library for building neural networks with named tensors, in the tradition of Alexander Rush's [Tensor Considered Harmful](https://nlp.seas.harvard.edu/NamedTensor).
 Named tensors improve the **legibility** and **compositionality** of tensor programs by using named axes instead of positional indices
 as typically used in NumPy, PyTorch, etc.
 

{haliax-1.4.dev381 → haliax-1.4.dev386}/README.md

@@ -10,14 +10,14 @@
 <a href="">
 <img alt="License" src="https://img.shields.io/github/license/stanford-crfm/haliax?color=blue" />
 </a>
-<a href="https://https://pypi.org/project/haliax/">
+<a href="https://pypi.org/project/haliax/">
     <img alt="PyPI" src="https://img.shields.io/pypi/v/haliax?color=blue" />
 </a>
 
 > *Though you don’t seem to be much for listening, it’s best to be careful. If you managed to catch hold of even just a piece of my name, you’d have all manner of power over me.*<br/>
 > — Patrick Rothfuss, *The Name of the Wind*
 
-Haliax is a [JAX](https:://github.com/google/jax) library for building neural networks with named tensors, in the tradition of Alexander Rush's [Tensor Considered Harmful](https://nlp.seas.harvard.edu/NamedTensor).
+Haliax is a [JAX](https://github.com/google/jax) library for building neural networks with named tensors, in the tradition of Alexander Rush's [Tensor Considered Harmful](https://nlp.seas.harvard.edu/NamedTensor).
 Named tensors improve the **legibility** and **compositionality** of tensor programs by using named axes instead of positional indices
 as typically used in NumPy, PyTorch, etc.
 

{haliax-1.4.dev381 → haliax-1.4.dev386}/docs/indexing.md

@@ -1,7 +1,7 @@
 # Indexing and Slicing
 
 Haliax supports Numpy-style indexing, including so-called [Advanced Indexing](https://numpy.org/doc/stable/user/basics.indexing.html#advanced-indexing),
-though the syntax is necessarily different. Most forms of indexing are supporting, except we don't support indexing with
+though the syntax is necessarily different. Most forms of indexing are supported, except we don't support indexing with
 booleans right now. (JAX doesn't support indexing with non-constant bool arrays anyway,
 so I don't think it's worth the effort to implement it in Haliax.)
 

{haliax-1.4.dev381 → haliax-1.4.dev386}/docs/nn.md

@@ -6,7 +6,7 @@
 Haliax provides a small number of neural network modules that are compatible with Equinox, though
 they naturally all use [haliax.NamedArray][]. (We welcome PRs for more modules! Nothing too exotic though.)
 
-The most interesting of these modules is [haliax.nn.Stacked][], which allows you to create homogenous "stacks"
+The most interesting of these modules is [haliax.nn.Stacked][], which allows you to create homogeneous "stacks"
 of the same module (e.g. transformer blocks), which is a common pattern in deep learning.
 
 ### Linear

{haliax-1.4.dev381 → haliax-1.4.dev386}/docs/state-dict.md

@@ -226,7 +226,7 @@ any Axis members to match the new shape.
 ::: haliax.state_dict.save_state_dict
 ::: haliax.state_dict.load_state_dict
 
-### Converting betweewn State Dicts and Modules
+### Converting between State Dicts and Modules
 
 ::: haliax.state_dict.from_state_dict
 ::: haliax.state_dict.to_state_dict

{haliax-1.4.dev381 → haliax-1.4.dev386}/docs/typing.md

@@ -1,4 +1,4 @@
-from haliax import NamedArrayfrom haliax import NamedArray
+from haliax import NamedArray
 
 # NamedArray Type Annotations
 

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

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

{haliax-1.4.dev381 → haliax-1.4.dev386}/src/haliax/__init__.py

@@ -66,7 +66,7 @@ from .core import (
 from .haxtyping import Named
 from .hof import fold, map, scan, vmap
 from .jax_utils import tree_checkpoint_name
-from .ops import clip, isclose, pad_left, pad, trace, tril, triu, where
+from .ops import clip, isclose, pad_left, pad, trace, tril, triu, unique, where
 from .partitioning import auto_sharded, axis_mapping, fsdp, named_jit, shard, shard_with_axis_mapping
 from .specialized_fns import top_k
 from .types import Scalar
@@ -1034,6 +1034,7 @@ __all__ = [
     "vmap",
     "trace",
     "where",
+    "unique",
     "clip",
     "tril",
     "triu",

{haliax-1.4.dev381 → haliax-1.4.dev386}/src/haliax/nn/conv.py

@@ -213,7 +213,7 @@ class Conv(_ConvBase):
         return x
 
     def _do_conv(self, inputs):
-        # _do_conv expects there ot be a single __batch__ dimension
+        # _do_conv expects there to be a single __batch__ dimension
         output_axes = _compute_output_axes(inputs, "__batch__", self.In, self.Out)
 
         batch_index = _index_of_name(inputs.axes, "__batch__")

{haliax-1.4.dev381 → haliax-1.4.dev386}/src/haliax/nn/linear.py

@@ -143,7 +143,7 @@ class MoELinear(eqx.Module):
     Experts: AxisSpec = eqx.field(static=True)
     In: Axis = eqx.field(static=True)
     Out: Axis = eqx.field(static=True)
-    # TODO: support quanitization for ragged_dot?
+    # TODO: support quantization for ragged_dot?
     # dot_general: DotGeneralOp = eqx.field(default_factory=DotGeneralOp.default)
 
     use_gmm: bool = eqx.field(static=True)

{haliax-1.4.dev381 → haliax-1.4.dev386}/src/haliax/ops.py

@@ -3,6 +3,9 @@ from typing import Mapping, Optional, Union
 
 import jax
 import jax.numpy as jnp
+from jaxtyping import ArrayLike
+
+import haliax
 
 from .axis import Axis, AxisSelector, axis_name
 from .core import NamedArray, NamedOrNumeric, broadcast_arrays, broadcast_arrays_and_return_axes, named
@@ -196,4 +199,147 @@ def raw_array_or_scalar(x: NamedOrNumeric):
     return x
 
 
-__all__ = ["trace", "where", "tril", "triu", "isclose", "pad_left", "pad", "clip"]
+@typing.overload
+def unique(
+    array: NamedArray, Unique: Axis, *, axis: AxisSelector | None = None, fill_value: ArrayLike | None = None
+) -> NamedArray:
+    ...
+
+
+@typing.overload
+def unique(
+    array: NamedArray,
+    Unique: Axis,
+    *,
+    return_index: typing.Literal[True],
+    axis: AxisSelector | None = None,
+    fill_value: ArrayLike | None = None,
+) -> tuple[NamedArray, NamedArray]:
+    ...
+
+
+@typing.overload
+def unique(
+    array: NamedArray,
+    Unique: Axis,
+    *,
+    return_inverse: typing.Literal[True],
+    axis: AxisSelector | None = None,
+    fill_value: ArrayLike | None = None,
+) -> tuple[NamedArray, NamedArray]:
+    ...
+
+
+@typing.overload
+def unique(
+    array: NamedArray,
+    Unique: Axis,
+    *,
+    return_counts: typing.Literal[True],
+    axis: AxisSelector | None = None,
+    fill_value: ArrayLike | None = None,
+) -> tuple[NamedArray, NamedArray]:
+    ...
+
+
+@typing.overload
+def unique(
+    array: NamedArray,
+    Unique: Axis,
+    *,
+    return_index: bool = False,
+    return_inverse: bool = False,
+    return_counts: bool = False,
+    axis: AxisSelector | None = None,
+    fill_value: ArrayLike | None = None,
+) -> NamedArray | tuple[NamedArray, ...]:
+    ...
+
+
+def unique(
+    array: NamedArray,
+    Unique: Axis,
+    *,
+    return_index: bool = False,
+    return_inverse: bool = False,
+    return_counts: bool = False,
+    axis: AxisSelector | None = None,
+    fill_value: ArrayLike | None = None,
+) -> NamedArray | tuple[NamedArray, ...]:
+    """
+    Like jnp.unique, but with named axes.
+
+    Args:
+        array: The input array.
+        Unique: The name of the axis that will be created to hold the unique values.
+        fill_value: The value to use for the fill_value argument of jnp.unique
+        axis: The axis along which to find unique values.
+        return_index: If True, return the indices of the unique values.
+        return_inverse: If True, return the indices of the input array that would reconstruct the unique values.
+    """
+    size = Unique.size
+
+    is_multireturn = return_index or return_inverse or return_counts
+
+    kwargs = dict(
+        size=size,
+        fill_value=fill_value,
+        return_index=return_index,
+        return_inverse=return_inverse,
+        return_counts=return_counts,
+    )
+
+    if axis is not None:
+        axis_index = array._lookup_indices(axis)
+        if axis_index is None:
+            raise ValueError(f"Axis {axis} not found in array. Available axes: {array.axes}")
+        out = jnp.unique(array.array, axis=axis_index, **kwargs)
+    else:
+        out = jnp.unique(array.array, **kwargs)
+
+    if is_multireturn:
+        unique = out[0]
+        next_index = 1
+        if return_index:
+            index = out[next_index]
+            next_index += 1
+        if return_inverse:
+            inverse = out[next_index]
+            next_index += 1
+        if return_counts:
+            counts = out[next_index]
+            next_index += 1
+    else:
+        unique = out
+
+    ret = []
+
+    if axis is not None:
+        out_axes = haliax.axis.replace_axis(array.axes, axis, Unique)
+    else:
+        out_axes = (Unique,)
+
+    unique_values = haliax.named(unique, out_axes)
+    if not is_multireturn:
+        return unique_values
+
+    ret.append(unique_values)
+
+    if return_index:
+        ret.append(haliax.named(index, Unique))
+
+    if return_inverse:
+        if axis is not None:
+            assert axis_index is not None
+            inverse = haliax.named(inverse, array.axes[axis_index])
+        else:
+            inverse = haliax.named(inverse, array.axes)
+        ret.append(inverse)
+
+    if return_counts:
+        ret.append(haliax.named(counts, Unique))
+
+    return tuple(ret)
+
+
+__all__ = ["trace", "where", "tril", "triu", "isclose", "pad_left", "pad", "clip", "unique"]

{haliax-1.4.dev381 → haliax-1.4.dev386}/tests/test_ops.py

@@ -250,3 +250,112 @@ def test_pad():
     assert padded.axes[0].size == Height.size + 3
     assert padded.axes[1].size == Width.size + 1
     assert jnp.all(expected == padded.array)
+
+
+def test_unique():
+    # named version of this test:
+    # >>> M = jnp.array([[1, 2],
+    # ...                [2, 3],
+    # ...                [1, 2]])
+    # >>> jnp.unique(M)
+    # Array([1, 2, 3], dtype=int32)
+
+    Height = Axis("Height", 3)
+    Width = Axis("Width", 2)
+
+    named1 = hax.named([[1, 2], [2, 3], [1, 2]], (Height, Width))
+
+    U = Axis("U", 3)
+
+    named2 = hax.unique(named1, U)
+
+    assert jnp.all(jnp.equal(named2.array, jnp.array([1, 2, 3])))
+
+    #     If you pass an ``axis`` keyword, you can find unique *slices* of the array along
+    #     that axis:
+    #
+    #     >>> jnp.unique(M, axis=0)
+    #     Array([[1, 2],
+    #            [2, 3]], dtype=int32)
+
+    U2 = Axis("U2", 2)
+    named3 = hax.unique(named1, U2, axis=Height)
+    assert jnp.all(jnp.equal(named3.array, jnp.array([[1, 2], [2, 3]])))
+
+    #     >>> x = jnp.array([3, 4, 1, 3, 1])
+    #     >>> values, indices = jnp.unique(x, return_index=True)
+    #     >>> print(values)
+    #     [1 3 4]
+    #     >>> print(indices)
+    #     [2 0 1]
+    #     >>> jnp.all(values == x[indices])
+    #     Array(True, dtype=bool)
+
+    x = hax.named([3, 4, 1, 3, 1], ("Height",))
+    U3 = Axis("U3", 3)
+    values, indices = hax.unique(x, U3, return_index=True)
+
+    assert jnp.all(jnp.equal(values.array, jnp.array([1, 3, 4])))
+    assert jnp.all(jnp.equal(indices.array, jnp.array([2, 0, 1])))
+
+    assert jnp.all(jnp.equal(values.array, x[{"Height": indices}].array))
+
+    #   If you set ``return_inverse=True``, then ``unique`` returns the indices within the
+    #     unique values for every entry in the input array:
+    #
+    #     >>> x = jnp.array([3, 4, 1, 3, 1])
+    #     >>> values, inverse = jnp.unique(x, return_inverse=True)
+    #     >>> print(values)
+    #     [1 3 4]
+    #     >>> print(inverse)
+    #     [1 2 0 1 0]
+    #     >>> jnp.all(values[inverse] == x)
+    #     Array(True, dtype=bool)
+
+    values, inverse = hax.unique(x, U3, return_inverse=True)
+
+    assert jnp.all(jnp.equal(values.array, jnp.array([1, 3, 4])))
+    assert jnp.all(jnp.equal(inverse.array, jnp.array([1, 2, 0, 1, 0])))
+
+    #     In multiple dimensions, the input can be reconstructed using
+    #     :func:`jax.numpy.take`:
+    #
+    #     >>> values, inverse = jnp.unique(M, axis=0, return_inverse=True)
+    #     >>> jnp.all(jnp.take(values, inverse, axis=0) == M)
+    #     Array(True, dtype=bool)
+    #
+
+    values, inverse = hax.unique(named1, U3, axis=Height, return_inverse=True)
+
+    assert jnp.all((values[{"U3": inverse}] == named1).array)
+
+    #     **Returning counts**
+    #  If you set ``return_counts=True``, then ``unique`` returns the number of occurrences
+    #     within the input for every unique value:
+    #
+    #     >>> x = jnp.array([3, 4, 1, 3, 1])
+    #     >>> values, counts = jnp.unique(x, return_counts=True)
+    #     >>> print(values)
+    #     [1 3 4]
+    #     >>> print(counts)
+    #     [2 2 1]
+    #
+    #     For multi-dimensional arrays, this also returns a 1D array of counts
+    #     indicating number of occurrences along the specified axis:
+    #
+    #     >>> values, counts = jnp.unique(M, axis=0, return_counts=True)
+    #     >>> print(values)
+    #     [[1 2]
+    #      [2 3]]
+    #     >>> print(counts)
+    #     [2 1]
+
+    values, counts = hax.unique(x, U3, return_counts=True)
+
+    assert jnp.all(jnp.equal(values.array, jnp.array([1, 3, 4])))
+    assert jnp.all(jnp.equal(counts.array, jnp.array([2, 2, 1])))
+
+    values, counts = hax.unique(named1, U2, axis=Height, return_counts=True)
+
+    assert jnp.all(jnp.equal(values.array, jnp.array([[1, 2], [2, 3]])))
+    assert jnp.all(jnp.equal(counts.array, jnp.array([2, 1])))

{haliax-1.4.dev381 → haliax-1.4.dev386}/tests/test_rearrange.py

@@ -293,7 +293,7 @@ def test_examples():
     r = einops_rearrange(z, "{B (H: h1 h) (W: w1 w) C} -> (B: B h1 w1) ... (C: C h w) ", h1=2, w1=2)
     assert r.axes == (Axis("B", B.size * 2 * 2), D, Axis("C", C.size * sH.size * sW.size))
     # unet attention reordering:
-    # postional: (qkv heads c) h w -> qkv heads c (h w)
+    # positional: (qkv heads c) h w -> qkv heads c (h w)
     # named: { (embed: qkv heads c) h w } -> qkv heads c (pos: h w)
     Embed = Axis("embed", 3 * 4 * C.size)
     attn = hax.random.randint(PRNGKey(0), (Embed, H, W), 0, 255)

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

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