haliax 1.4.dev405__tar.gz → 1.4.dev407__tar.gz

haliax-1.4.dev407/.agents/projects/api_parity.md

@@ -0,0 +1,215 @@
+# JAX NumPy API Parity
+This document tracks JAX NumPy functions not yet wrapped by Haliax.
+
+APIs that don't translate well to named tensors are intentionally omitted here. This includes dtype constructors, raw array converters (e.g. `from_dlpack`), indexing helpers like `c_`/`r_`, and functions whose JAX counterparts already work with `NamedArray` out of the box. Basic array construction is handled by `haliax.named`, so functions such as `array` or `asarray` are not listed.
+
+## numpy
+- [x] `allclose`
+- [ ] `amin`
+- [ ] `append`
+- [ ] `apply_along_axis`
+- [ ] `apply_over_axes`
+- [ ] `argpartition`
+- [ ] `argwhere`
+- [x] `array_equal`
+- [x] `array_equiv`
+- [ ] `array_split`
+- [ ] `astype`
+- [ ] `atan2`
+- [ ] `average`
+- [ ] `bartlett`
+- [ ] `bitwise_count`
+- [ ] `bitwise_invert`
+- [ ] `bitwise_left_shift`
+- [ ] `bitwise_right_shift`
+- [ ] `blackman`
+- [ ] `block`
+- [ ] `broadcast_shapes`
+- [ ] `can_cast`
+- [ ] `choose`
+- [ ] `column_stack`
+- [ ] `compress`
+- [ ] `concat`
+- [ ] `convolve`
+- [ ] `copysign`
+- [ ] `corrcoef`
+- [ ] `correlate`
+- [ ] `count_nonzero`
+- [ ] `cov`
+- [ ] `cumulative_prod`
+- [ ] `cumulative_sum`
+- [ ] `delete`
+- [ ] `diag`
+- [ ] `diag_indices`
+- [ ] `diag_indices_from`
+- [ ] `diagflat`
+- [ ] `diagonal`
+- [ ] `diff`
+- [ ] `digitize`
+- [ ] `dsplit`
+- [ ] `dstack`
+- [ ] `ediff1d`
+- [ ] `einsum_path`
+- [ ] `empty`
+- [ ] `empty_like`
+- [ ] `expand_dims`
+- [ ] `extract`
+- [ ] `eye`
+- [ ] `fill_diagonal`
+- [ ] `flatnonzero`
+- [ ] `flip`
+- [ ] `fliplr`
+- [ ] `flipud`
+- [ ] `gcd`
+- [ ] `gradient`
+- [ ] `hamming`
+- [ ] `hanning`
+- [ ] `heaviside`
+- [ ] `histogram`
+- [ ] `histogram2d`
+- [ ] `histogram_bin_edges`
+- [ ] `histogramdd`
+- [ ] `hsplit`
+- [ ] `hstack`
+- [ ] `identity`
+- [ ] `indices`
+- [ ] `inner`
+- [ ] `insert`
+- [ ] `interp`
+- [ ] `intersect1d`
+- [ ] `iscomplexobj`
+- [ ] `isin`
+- [ ] `iterable`
+- [ ] `ix_`
+- [ ] `kaiser`
+- [ ] `kron`
+- [ ] `lcm`
+- [ ] `ldexp`
+- [ ] `lexsort`
+- [ ] `mask_indices`
+- [ ] `median`
+- [ ] `meshgrid`
+- [ ] `mgrid`
+- [ ] `modf`
+- [ ] `moveaxis`
+- [ ] `nan_to_num`
+- [ ] `nanargmax`
+- [ ] `nanargmin`
+- [ ] `nancumprod`
+- [ ] `nancumsum`
+- [ ] `nanmax`
+- [ ] `nanmean`
+- [ ] `nanmedian`
+- [ ] `nanmin`
+- [ ] `nanpercentile`
+- [ ] `nanprod`
+- [ ] `nanquantile`
+- [ ] `nanstd`
+- [ ] `nansum`
+- [ ] `nanvar`
+- [ ] `nonzero`
+- [ ] `ogrid`
+- [ ] `packbits`
+- [ ] `partition`
+- [ ] `percentile`
+- [ ] `permute_dims`
+- [ ] `piecewise`
+- [ ] `place`
+- [ ] `poly`
+- [ ] `polyadd`
+- [ ] `polyder`
+- [ ] `polydiv`
+- [ ] `polyfit`
+- [ ] `polyint`
+- [ ] `polymul`
+- [ ] `polysub`
+- [ ] `polyval`
+- [ ] `pow`
+- [ ] `promote_types`
+- [ ] `put`
+- [ ] `put_along_axis`
+- [ ] `quantile`
+- [ ] `ravel_multi_index`
+- [ ] `reshape`
+- [ ] `resize`
+- [ ] `result_type`
+- [ ] `rollaxis`
+- [ ] `roots`
+- [ ] `rot90`
+- [ ] `select`
+- [ ] `setdiff1d`
+- [ ] `setxor1d`
+- [ ] `spacing`
+- [ ] `squeeze`
+- [ ] `swapaxes`
+- [ ] `take_along_axis`
+- [ ] `transpose`
+- [ ] `trapezoid`
+- [ ] `tri`
+- [ ] `tril_indices`
+- [ ] `tril_indices_from`
+- [ ] `trim_zeros`
+- [ ] `triu_indices`
+- [ ] `triu_indices_from`
+- [ ] `union1d`
+- [ ] `unpackbits`
+- [ ] `unravel_index`
+- [ ] `unstack`
+- [ ] `unwrap`
+- [ ] `vander`
+- [ ] `vsplit`
+- [ ] `vstack`
+
+## fft
+- [ ] `fft`
+- [ ] `fft2`
+- [ ] `fftfreq`
+- [ ] `fftn`
+- [ ] `fftshift`
+- [ ] `hfft`
+- [ ] `ifft`
+- [ ] `ifft2`
+- [ ] `ifftn`
+- [ ] `ifftshift`
+- [ ] `ihfft`
+- [ ] `irfft`
+- [ ] `irfft2`
+- [ ] `irfftn`
+- [ ] `rfft`
+- [ ] `rfft2`
+- [ ] `rfftfreq`
+- [ ] `rfftn`
+
+## linalg
+- [ ] `cholesky`
+- [ ] `cond`
+- [ ] `cross`
+- [ ] `det`
+- [ ] `diagonal`
+- [ ] `eig`
+- [ ] `eigh`
+- [ ] `eigvals`
+- [ ] `eigvalsh`
+- [ ] `inv`
+- [ ] `lstsq`
+- [ ] `matmul`
+- [ ] `matrix_norm`
+- [ ] `matrix_power`
+- [ ] `matrix_rank`
+- [ ] `matrix_transpose`
+- [ ] `norm`
+- [ ] `outer`
+- [ ] `pinv`
+- [ ] `qr`
+- [ ] `slogdet`
+- [ ] `solve`
+- [ ] `svd`
+- [ ] `svdvals`
+- [ ] `tensordot`
+- [ ] `tensorinv`
+- [ ] `tensorsolve`
+- [ ] `trace`
+- [ ] `vdot`
+- [ ] `vecdot`
+- [ ] `vecmat`
+- [ ] `vector_norm`

{haliax-1.4.dev405 → haliax-1.4.dev407}/PKG-INFO

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

@@ -259,6 +259,9 @@ These are all more or less directly from JAX's NumPy API.
 ::: haliax.bincount
 ::: haliax.clip
 ::: haliax.isclose
+::: haliax.allclose
+::: haliax.array_equal
+::: haliax.array_equiv
 ::: haliax.pad
 ::: haliax.searchsorted
 ::: haliax.top_k

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

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

{haliax-1.4.dev405 → haliax-1.4.dev407}/src/haliax/__init__.py

@@ -69,6 +69,9 @@ from .hof import fold, map, scan, vmap
 from .jax_utils import tree_checkpoint_name
 from .ops import (
     clip,
+    allclose,
+    array_equal,
+    array_equiv,
     isclose,
     pad_left,
     pad,
@@ -1105,6 +1108,9 @@ __all__ = [
     "shard",
     "enable_shape_checks",
     "are_shape_checks_enabled",
+    "allclose",
+    "array_equal",
+    "array_equiv",
     "isclose",
     "pad_left",
     "pad",

{haliax-1.4.dev405 → haliax-1.4.dev407}/src/haliax/ops.py

@@ -138,6 +138,29 @@ def isclose(a: NamedArray, b: NamedArray, rtol=1e-05, atol=1e-08, equal_nan=Fals
     return NamedArray(jnp.isclose(a.array, b.array, rtol=rtol, atol=atol, equal_nan=equal_nan), a.axes)
 
 
+def allclose(a: NamedArray, b: NamedArray, rtol=1e-05, atol=1e-08, equal_nan=False) -> bool:
+    """Returns True if two arrays are element-wise equal within a tolerance."""
+    a, b = broadcast_arrays(a, b)
+    return bool(jnp.allclose(a.array, b.array, rtol=rtol, atol=atol, equal_nan=equal_nan))
+
+
+def array_equal(a: NamedArray, b: NamedArray) -> bool:
+    """Returns True if two arrays have the same shape and elements."""
+    if set(a.axes) != set(b.axes):
+        return False
+    b = b.rearrange(a.axes)
+    return bool(jnp.array_equal(a.array, b.array))
+
+
+def array_equiv(a: NamedArray, b: NamedArray) -> bool:
+    """Returns True if two arrays are shape-consistent and equal."""
+    try:
+        a, b = broadcast_arrays(a, b)
+    except ValueError:
+        return False
+    return bool(jnp.array_equal(a.array, b.array))
+
+
 def pad_left(array: NamedArray, axis: Axis, new_axis: Axis, value=0) -> NamedArray:
     """Pad an array along named axes."""
     amount_to_pad_to = new_axis.size - axis.size

{haliax-1.4.dev405 → haliax-1.4.dev407}/tests/test_ops.py

@@ -425,6 +425,32 @@ def test_bincount():
     assert jnp.allclose(out_w.array, expected_w)
 
 
+def test_allclose_array_equal_equiv():
+    A = Axis("A", 2)
+    B = Axis("B", 3)
+    x = hax.random.uniform(PRNGKey(0), (A, B))
+    y = x + 1e-6
+
+    assert hax.allclose(x, y)
+    assert not hax.allclose(x, x + 1.0)
+
+    x1 = hax.ones((A, B))
+    y_reordered = x1.rearrange((B, A))
+    assert hax.array_equal(x1, y_reordered)
+
+    scalar = hax.ones(())
+    assert hax.array_equiv(x1, scalar)
+    assert not hax.array_equal(x1, scalar)
+
+    y_vec = hax.ones((B,))
+    assert hax.array_equiv(x1, y_vec)
+    assert not hax.array_equal(x1, y_vec)
+
+    C = Axis("C", 4)
+    z = hax.ones((C,))
+    assert not hax.array_equiv(x1, z)
+
+
 def test_roll_scalar_named_shift():
     H = Axis("H", 4)
     W = Axis("W", 3)

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

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