haliax 1.4.dev379__tar.gz → 1.4.dev380__tar.gz

{haliax-1.4.dev379 → haliax-1.4.dev380}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: haliax
-Version: 1.4.dev379
+Version: 1.4.dev380
 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.dev379 → haliax-1.4.dev380}/docs/api.md

@@ -258,6 +258,7 @@ These are all more or less directly from JAX's NumPy API.
 
 ::: haliax.clip
 ::: haliax.isclose
+::: haliax.pad
 ::: haliax.top_k
 ::: haliax.trace
 ::: haliax.tril

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

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

{haliax-1.4.dev379 → haliax-1.4.dev380}/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, trace, tril, triu, where
+from .ops import clip, isclose, pad_left, pad, trace, tril, triu, 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
@@ -1082,6 +1082,7 @@ __all__ = [
     "are_shape_checks_enabled",
     "isclose",
     "pad_left",
+    "pad",
     "stack",
     "concatenate",
     "eliminate_axes",

{haliax-1.4.dev379 → haliax-1.4.dev380}/src/haliax/ops.py

@@ -1,10 +1,10 @@
 import typing
-from typing import Optional, Union
+from typing import Mapping, Optional, Union
 
 import jax
 import jax.numpy as jnp
 
-from .axis import Axis, AxisSelector
+from .axis import Axis, AxisSelector, axis_name
 from .core import NamedArray, NamedOrNumeric, broadcast_arrays, broadcast_arrays_and_return_axes, named
 from .jax_utils import is_scalarish
 
@@ -152,10 +152,48 @@ def pad_left(array: NamedArray, axis: Axis, new_axis: Axis, value=0) -> NamedArr
     return NamedArray(padded, array.axes[:idx] + (new_axis,) + array.axes[idx + 1 :])
 
 
+def pad(
+    array: NamedArray,
+    pad_width: Mapping[AxisSelector, tuple[int, int]],
+    *,
+    mode: str = "constant",
+    constant_values: NamedOrNumeric = 0,
+    **kwargs,
+) -> NamedArray:
+    """Version of ``jax.numpy.pad`` that works with ``NamedArray``.
+
+    ``pad_width`` should be a mapping from axis (or axis name) to a ``(before, after)``
+    tuple specifying how much padding to add on each side of that axis. Any axis
+    not present in ``pad_width`` will not be padded.
+    """
+
+    padding = []
+    new_axes = []
+    for ax in array.axes:
+        left_right = pad_width.get(ax)
+        if left_right is None:
+            left_right = pad_width.get(axis_name(ax))  # type: ignore[arg-type]
+        if left_right is None:
+            left_right = (0, 0)
+        left, right = left_right
+        padding.append((left, right))
+        new_axes.append(ax.resize(ax.size + left + right))
+
+    result = jnp.pad(
+        array.array,
+        padding,
+        mode=mode,
+        constant_values=raw_array_or_scalar(constant_values),
+        **kwargs,
+    )
+
+    return NamedArray(result, tuple(new_axes))
+
+
 def raw_array_or_scalar(x: NamedOrNumeric):
     if isinstance(x, NamedArray):
         return x.array
     return x
 
 
-__all__ = ["trace", "where", "tril", "triu", "isclose", "pad_left", "clip"]
+__all__ = ["trace", "where", "tril", "triu", "isclose", "pad_left", "pad", "clip"]

{haliax-1.4.dev379 → haliax-1.4.dev380}/tests/test_ops.py

@@ -237,3 +237,16 @@ def test_reductions_produce_scalar_named_arrays_when_None_axis():
     # But if we specify axes, we always get a NamedArray, even if it's a scalar
     assert isinstance(hax.mean(named1, axis=("Height", "Width")), NamedArray)
     assert hax.mean(named1, axis=("Height", "Width")).axes == ()
+
+
+def test_pad():
+    Height = Axis("Height", 3)
+    Width = Axis("Width", 2)
+
+    arr = hax.arange((Height, Width))
+    padded = hax.pad(arr, {Height: (1, 2), Width: (0, 1)}, mode="constant", constant_values=0)
+
+    expected = jnp.pad(arr.array, [(1, 2), (0, 1)], mode="constant", constant_values=0)
+    assert padded.axes[0].size == Height.size + 3
+    assert padded.axes[1].size == Width.size + 1
+    assert jnp.all(expected == padded.array)

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

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