# Copyright 2026 The JAX Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Tests for experimental searchsorted primitive."""

import functools

from absl.testing import absltest
import jax
import jax.numpy as jnp
import numpy as np

from jax._src import config
from jax._src.numpy import hijax
from jax._src import test_util as jtu
from jax.sharding import PartitionSpec as P

jtu.request_cpu_devices(4)

config.parse_flags_with_absl()

_SIDES = ("left", "right")
_DTYPES = ("int32", "float32", "complex64")
_METHODS = ("compare_all", "scan", "scan_unrolled", "sort")
_QUERY_SHAPES = ((), (5,), (5, 7))
_SORTED_ARR_SHAPES = ((10,), (10, 11), (10, 11, 12))
_BATCH_SHAPES = ((), (3,), (3, 4))

searchsorted_jit = jax.jit(
    hijax.searchsorted,
    static_argnames=["side", "dimension", "batch_dims", "method"],
)


def searchsorted_reference(
    sorted_arr: np.ndarray,
    query: np.ndarray,
    *,
    side: str = "left",
    method: str | None = None,
    dimension: int = 0,
    batch_dims: int = 0,
    dtype: np.dtype = np.int32,
) -> np.ndarray:
  """Reference implementation of the searchsorted primitive in terms of numpy.

  Args:
    sorted_arr: N-dimensional array, which is assumed to be sorted in increasing
      order along ``dimension``.
    query: M-dimensional array of query values.
    side: {'left', 'right'}. If 'left', find the index of the first suitable
      location. If 'right', find the index of the last.
    method: unused.
    dimension: positive integer specifying the dimension along which to insert
      query values.
    batch_dims: integer specifying the number of leading dimensions of
      `sorted_arr` and `query` to treat as batch dimensions.

  Returns:
    An array specifying the insertion locations of `query` into `sorted_arr`.
  """
  del method  # unused in numpy.searchsorted.

  if batch_dims < 0 or batch_dims >= sorted_arr.ndim:
    raise ValueError(
        f"batch_dims={batch_dims} must be in range [0, {sorted_arr.ndim})"
    )

  if dimension < batch_dims or dimension >= sorted_arr.ndim:
    raise ValueError(
        f"dimension={dimension} must be in range [{batch_dims},"
        f" {sorted_arr.ndim})"
    )

  if sorted_arr.shape[:batch_dims] != query.shape[:batch_dims]:
    raise ValueError(
        "batch dimension sizes must match; got"
        f" {sorted_arr.shape[:batch_dims]} != {query.shape[:batch_dims]}"
    )

  # Recursive call to handle common batch dimensions.
  if batch_dims > 0:
    return np.array([
        searchsorted_reference(
            a, q, side=side, dimension=dimension - 1, batch_dims=batch_dims - 1
        )
        for a, q in zip(sorted_arr, query)
    ])

  # Recursive call to handle sorted array batch dimensions.
  if sorted_arr.ndim > 1:
    sorted_arr = np.moveaxis(sorted_arr, dimension, 1)
    return np.array(
        [searchsorted_reference(a, query, side=side) for a in sorted_arr]
    )

  # Base case: batched query handled by numpy.
  assert sorted_arr.ndim == 1

  return np.searchsorted(sorted_arr, query, side=side).astype(dtype)

def nonzero_reference(
    a: np.ndarray,
    size: int, *,
    fill_value: int | tuple[int, ...] | None = None,
    axes: tuple[int, ...] | None = None,
    dtype: np.dtype | str = "int32",
) -> tuple[np.ndarray, ...]:
  """Reference implementation for hijax.nonzero in terms of np.nonzero.

  Args:
    a: n-dimensional array
    size: integer specifying size of output
    fill_value: optional padding value. Defaults to 0.
    axes: optional set of axes over which to compute the indices.
      Defaults to range(a.ndim).
    dtype: optional datatype for the returned indices. Defaults to int32.

  Returns:
    A tuple of indices corresponding to axes
  """
  a = np.asarray(a)
  if axes is None:
    axes = tuple(range(a.ndim))
  else:
    axes = tuple(sorted(ax % a.ndim for ax in axes))

  if fill_value is not None:
    if isinstance(fill_value, tuple):
      assert len(fill_value) == len(axes)
      fill_value_tup = fill_value
    else:
      fill_value_tup = (fill_value,) * len(axes)
  else:
    fill_value_tup = (0,) * len(axes)

  batch_axes = [ax for ax in range(a.ndim) if ax not in axes]
  transposed_a = np.transpose(a, (*batch_axes, *axes))
  batch_dims = transposed_a.shape[:len(batch_axes)]

  res = []
  for idx in np.ndindex(batch_dims):
    sub_a = transposed_a[idx]
    nz = np.nonzero(sub_a)

    true_size = len(nz[0])
    nz_padded = []
    for axis_idx, axis_nz in enumerate(nz):
      if true_size >= size:
        nz_padded.append(axis_nz[:size])
      else:
        # Pad with fill_value
        nz_padded.append(np.pad(axis_nz, (0, size - true_size), mode="constant",
                                constant_values=fill_value_tup[axis_idx]))
    res.append(nz_padded)

  out = []
  for axis_idx in range(len(axes)):
    axis_res = [r[axis_idx] for r in res]
    if batch_dims:
      axis_res_array = np.array(axis_res).reshape(*batch_dims, size)
    else:
      axis_res_array = np.array(axis_res).reshape(size)
    out.append(axis_res_array.astype(dtype))

  return tuple(out)


class SearchsortedTest(jtu.JaxTestCase):

  @jtu.sample_product(
      dtype=_DTYPES,
      out_dtype=['int32', 'uint32']
  )
  def test_out_dtype(self, dtype, out_dtype):
    rand_dtype = jtu.rand_default(self.rng())
    sorted_arr = np.sort(rand_dtype((100,), dtype))
    query = rand_dtype((100,), dtype)
    result = hijax.searchsorted(sorted_arr, query, dtype=out_dtype)
    self.assertEqual(result.dtype, np.dtype(out_dtype))

  @jtu.sample_product(
      dtype=_DTYPES,
      side=_SIDES,
      method=_METHODS,
  )
  def test_1D_against_numpy(self, dtype, side, method):
    rand_dtype = jtu.rand_default(self.rng())
    sorted_arr = np.sort(rand_dtype((100,), dtype))
    query = rand_dtype((100,), dtype)
    expected = np.searchsorted(sorted_arr, query, side=side)

    with self.subTest(name="non-jit"):
      actual = hijax.searchsorted(
          sorted_arr, query, side=side, method=method
      )
      np.testing.assert_array_equal(expected, actual)

    with self.subTest(name="jit"):
      actual = searchsorted_jit(sorted_arr, query, side=side, method=method)
      np.testing.assert_array_equal(expected, actual)

  @jtu.sample_product(
      side=_SIDES,
      method=_METHODS,
  )
  def test_2D_single_query_against_numpy(self, side, method):
    sorted_arr = np.sort(
        self.rng().uniform(size=(10, 100), low=0, high=10).round(1), axis=1
    )
    query = self.rng().uniform(size=(50,), low=-0.5, high=10.5).round(1)
    expected = np.array(
        [np.searchsorted(arr, query, side=side) for arr in sorted_arr]
    )

    with self.subTest(name="non-jit"):
      actual = hijax.searchsorted(
          sorted_arr, query, dimension=1, side=side, method=method
      )
      np.testing.assert_array_equal(expected, actual)

    with self.subTest(name="jit"):
      actual = searchsorted_jit(
          sorted_arr, query, dimension=1, side=side, method=method
      )
      np.testing.assert_array_equal(expected, actual)

  @jtu.sample_product(
      side=_SIDES,
      method=_METHODS,
  )
  def test_2D_batched_query_against_numpy(self, side, method):
    sorted_arr = np.sort(
        self.rng().uniform(size=(10, 100), low=0, high=10).round(1), axis=1
    )
    query = self.rng().uniform(size=(10, 20, 30), low=-0.5, high=10.5).round(1)
    expected = np.array(
        [np.searchsorted(a, q, side=side) for a, q in zip(sorted_arr, query)]
    )

    with self.subTest(name="non-jit"):
      actual = hijax.searchsorted(
          sorted_arr, query, dimension=1, batch_dims=1, side=side, method=method
      )
      np.testing.assert_array_equal(expected, actual)

    with self.subTest(name="jit"):
      actual = searchsorted_jit(
          sorted_arr, query, dimension=1, batch_dims=1, side=side, method=method
      )
      np.testing.assert_array_equal(expected, actual)

  @jtu.sample_product(
      side=_SIDES,
      method=_METHODS,
      query_shape=_QUERY_SHAPES,
      sorted_arr_shape=_SORTED_ARR_SHAPES,
      batch_shape=_BATCH_SHAPES,
  )
  def test_all_shapes(
      self, side, method, query_shape, sorted_arr_shape, batch_shape
  ):
    rand_dtype = jtu.rand_default(self.rng())
    batch_dims = len(batch_shape)
    dimension = batch_dims + int(self.rng().randint(0, len(sorted_arr_shape)))
    sorted_arr = np.sort(
        rand_dtype((*batch_shape, *sorted_arr_shape), np.float32),
        axis=dimension,
    )
    query = rand_dtype((*batch_shape, *query_shape), np.float32)
    kwds = dict(
        side=side,
        dimension=dimension,
        batch_dims=batch_dims,
        method=method,
    )
    expected = searchsorted_reference(sorted_arr, query, **kwds)
    actual = hijax.searchsorted(sorted_arr, query, **kwds)
    actual_jit = searchsorted_jit(sorted_arr, query, **kwds)

    with self.subTest(name="non-jit"):
      np.testing.assert_array_equal(expected, actual)

    with self.subTest(name="jit"):
      np.testing.assert_array_equal(expected, actual_jit)

  def test_input_validation(self):
    with self.subTest(name="invalid_side"):
      sorted_arr = np.arange(10)
      query = np.arange(10)
      with self.assertRaisesWithLiteralMatch(
          ValueError, "invalid argument side='none', expected 'left' or 'right'"
      ):
        hijax.searchsorted(sorted_arr, query, side="none")

    with self.subTest(name="invalid_batch_dims"):
      sorted_arr = np.arange(10)
      query = np.arange(10)
      with self.assertRaisesWithLiteralMatch(
          ValueError, "batch_dims=-1 must be in range [0, 1)"
      ):
        hijax.searchsorted(sorted_arr, query, batch_dims=-1)

    with self.subTest(name="invalid_dimension"):
      sorted_arr = np.arange(10)
      query = np.arange(10)
      with self.assertRaisesWithLiteralMatch(
          ValueError, "dimension=2 must be in range [0, 1)"
      ):
        hijax.searchsorted(sorted_arr, query, dimension=2)

    with self.subTest(name="batch_dimension_size_mismatch"):
      sorted_arr = np.zeros((2, 5))
      query = np.zeros((3, 6))
      with self.assertRaisesWithLiteralMatch(
          ValueError, "batch dimension sizes must match; got (2,) != (3,)"
      ):
        hijax.searchsorted(sorted_arr, query, batch_dims=1, dimension=1)

    with self.subTest(name="dtype_mismatch"):
      sorted_arr = np.zeros(5, dtype="float32")
      query = np.zeros(5, dtype="int32")
      with self.assertRaisesWithLiteralMatch(
          ValueError,
          "dtypes of sorted_arr and query must match; got float32 and int32",
      ):
        hijax.searchsorted(sorted_arr, query)

    with self.subTest(name="invalid_method"):
      sorted_arr = np.arange(5)
      query = 2
      with self.assertRaisesWithLiteralMatch(
          ValueError,
          "invalid argument method='none', expected one of ['compare_all',"
          " 'scan', 'scan_unrolled', 'sort']",
      ):
        hijax.searchsorted(sorted_arr, query, method="none")

  @jtu.sample_product(
      side=_SIDES,
      method=_METHODS,
  )
  def test_vmap(self, side, method):
    with self.subTest(name="batched sorted_arr"):
      sorted_arr = np.sort(
          self.rng().uniform(size=(10, 100), low=0, high=10).round(1), axis=1
      )
      query = self.rng().uniform(size=50, low=-0.5, high=10.5).round(1)
      expected = np.array(
          [np.searchsorted(a, query, side=side) for a in sorted_arr]
      )
      actual = jax.vmap(
          functools.partial(
              hijax.searchsorted, side=side, method=method
          ),
          in_axes=(0, None),
          out_axes=0,
      )(sorted_arr, query)
      np.testing.assert_array_equal(expected, actual)

    with self.subTest(name="batched query"):
      sorted_arr = np.sort(self.rng().uniform(size=100, low=0, high=10).round(1))
      query = self.rng().uniform(size=(10, 50), low=-0.5, high=10.5).round(1)
      expected = np.array(
          [np.searchsorted(sorted_arr, q, side=side) for q in query]
      )
      actual = jax.vmap(
          functools.partial(
              hijax.searchsorted, side=side, method=method
          ),
          in_axes=(None, 0),
          out_axes=0,
      )(sorted_arr, query)
      np.testing.assert_array_equal(expected, actual)

    with self.subTest(name="batched sorted_arr and query"):
      sorted_arr = np.sort(
          self.rng().uniform(size=(10, 100), low=0, high=10).round(1), axis=1
      )
      query = self.rng().uniform(size=(10, 50), low=-0.5, high=10.5).round(1)
      expected = np.array(
          [np.searchsorted(a, q, side=side) for a, q in zip(sorted_arr, query)]
      )
      actual = jax.vmap(
          functools.partial(
              hijax.searchsorted,
              side=side,
              method=method,
          )
      )(sorted_arr, query)
      np.testing.assert_array_equal(expected, actual)

  def test_vmap_unbatched(self):
    """Test batched case where neither input is batched"""
    x = np.array([2, 1, 3], dtype='int32')
    y = np.array([2, 3, 4, 5])
    z = np.array([4, 3])
    def f(x, y, z):
      return x * hijax.searchsorted(y, z)
    result = jax.vmap(f, (0, None, None))(x, y, z)
    expected = x[:, None] * hijax.searchsorted(y, z)[None, :]
    self.assertArraysEqual(result, expected)

  def test_nested_vmap(self):
    x = np.arange(60).reshape(2, 3, 10)
    y = np.arange(5, 65, 10).reshape(2, 3)
    out = jax.vmap(jax.vmap(hijax.searchsorted))(x, y)
    expected = jax.numpy.array([
      [hijax.searchsorted(x, y) for x, y in zip(xrow, yrow)]
      for xrow, yrow in zip(x, y)
    ])
    self.assertArraysEqual(out, expected)

  @jtu.sample_product(
      side=_SIDES,
      method=_METHODS,
  )
  def test_autodiff(self, side, method, dtype="float32"):
    rand_dtype = jtu.rand_default(self.rng())
    sorted_arr = np.sort(rand_dtype((100,), dtype))
    query = rand_dtype((100,), dtype)
    func = functools.partial(
        hijax.searchsorted, side=side, method=method
    )
    primal_expected = func(sorted_arr, query)

    with self.subTest(name="forward mode"):
      primals = (sorted_arr, query)
      tangents = (np.ones_like(sorted_arr), np.ones_like(query))
      primal_out, tangent_out = jax.jvp(func, primals, tangents)
      np.testing.assert_array_equal(primal_expected, primal_out, strict=True)
      self.assertEqual(tangent_out.shape, primal_out.shape)
      self.assertEqual(tangent_out.dtype, jax.dtypes.float0)

    with self.subTest(name="reverse mode"):
      primal_out, f_vjp = jax.vjp(func, sorted_arr, query)
      tangent_out = np.ones_like(primal_out)
      tangents = f_vjp(tangent_out)
      np.testing.assert_array_equal(primal_expected, primal_out, strict=True)
      np.testing.assert_array_equal(
          tangents[0], np.zeros_like(sorted_arr), strict=True
      )
      np.testing.assert_array_equal(
          tangents[1], np.zeros_like(query), strict=True
      )

  def test_reference_batch_dims_out_of_range(self):
    sorted_arr = np.arange(10)
    query = np.arange(10)
    with self.assertRaisesWithLiteralMatch(
        ValueError, "batch_dims=2 must be in range [0, 1)"
    ):
      searchsorted_reference(sorted_arr, query, batch_dims=2)

  def test_reference_dimension_out_of_range(self):
    sorted_arr = np.arange(10)
    query = np.arange(10)
    with self.assertRaisesWithLiteralMatch(
        ValueError, "dimension=2 must be in range [0, 1)"
    ):
      searchsorted_reference(sorted_arr, query, dimension=2)

  def test_reference_batch_dimension_size_mismatch(self):
    sorted_arr = np.zeros((2, 5))
    query = np.zeros((3, 6))
    with self.assertRaisesWithLiteralMatch(
        ValueError, "batch dimension sizes must match; got (2,) != (3,)"
    ):
      searchsorted_reference(
          sorted_arr, query, batch_dims=1, dimension=1
      )


class NonzeroTest(jtu.JaxTestCase):

  def test_1D_against_numpy(self):
    a = self.rng().randint(0, 2, size=(100,))
    expected = tuple(i.astype('int32') for i in np.nonzero(a))
    actual = hijax.nonzero(a, size=np.sum(a != 0))

    self.assertEqual(len(expected), len(actual))
    for e, a_i in zip(expected, actual):
      self.assertArraysEqual(e, a_i)

  @jtu.sample_product(
      axes=[None, (0,), (1,), (0, 1)],
  )
  def test_2D_against_reference(self, axes):
    a = self.rng().randint(0, 2, size=(10, 20))
    size = 5
    expected = nonzero_reference(a, size=size, axes=axes)
    actual = hijax.nonzero(a, size=size, axes=axes)

    self.assertEqual(len(expected), len(actual))
    for e, a_i in zip(expected, actual):
      self.assertArraysEqual(e, a_i)

  def test_vmap(self):
    a = self.rng().randint(0, 2, size=(3, 4, 5))
    size = 2

    def f(x):
      return hijax.nonzero(x, size=size, axes=(1,))[0]

    vmapped = jax.vmap(f)(a)

    expected = []
    for sub_a in a:
      expected.append(nonzero_reference(sub_a, size=size, axes=(1,))[0])
    expected = np.array(expected)

    self.assertArraysEqual(expected, vmapped)

  def test_jit(self):
    a = self.rng().randint(0, 2, size=(10,))
    size = 2

    @jax.jit
    def f(x):
      return hijax.nonzero(x, size=size)

    actual = f(a)
    expected = nonzero_reference(a, size=size)

    for e, a_i in zip(expected, actual):
      self.assertArraysEqual(e, a_i)

  def test_autodiff(self):
    a = jax.numpy.array([0.0, 1.0, 0.0, 1.0])
    size = 2

    def f(x):
      return hijax.nonzero(x, size=size)[0]

    primal_out, tangent_out = jax.jvp(f, (a,), (jax.numpy.ones_like(a),))
    self.assertEqual(tangent_out.dtype, jax.dtypes.float0)

    _, f_vjp = jax.vjp(f, a)
    tangent_in = f_vjp(jax.numpy.ones_like(primal_out))
    self.assertArraysEqual(tangent_in[0], jax.numpy.zeros_like(a))

  def test_1D_fill_value(self):
    a = self.rng().randint(0, 2, size=(100,))
    size = 150
    fill_val = 99
    expected = nonzero_reference(a, size=size, fill_value=fill_val)
    actual = hijax.nonzero(a, size=size, fill_value=fill_val)
    self.assertEqual(len(expected), len(actual))
    for e, a_i in zip(expected, actual):
      self.assertArraysEqual(e, a_i)

  @jtu.sample_product(
      axes=[None, (0,), (1,), (0, 1)],
      fill_value=[99, (99, 98)],
  )
  def test_2D_fill_value(self, axes, fill_value):
    a = self.rng().randint(0, 2, size=(10, 20))
    size = 250
    actual_axes = tuple(range(a.ndim)) if axes is None else axes
    if isinstance(fill_value, tuple) and len(fill_value) != len(actual_axes):
      self.skipTest("fill_value length must match number of axes")

    expected = nonzero_reference(a, size=size, fill_value=fill_value, axes=axes)
    actual = hijax.nonzero(a, size=size, fill_value=fill_value, axes=axes)

    self.assertEqual(len(expected), len(actual))
    for e, a_i in zip(expected, actual):
      self.assertArraysEqual(e, a_i)

  def test_vmap_fill_value_batched(self):
    a = self.rng().randint(0, 2, size=(3, 4, 5))
    fill_val = jax.numpy.array([99, 98, 97])
    size = 20

    vmapped = jax.vmap(lambda x, fv: hijax.nonzero(x, size=size, fill_value=fv, axes=(1,))[0])(a, fill_val)

    expected = []
    for sub_a, fv in zip(a, fill_val):
      expected.append(nonzero_reference(sub_a, size=size, fill_value=int(fv), axes=(1,))[0])
    expected = np.array(expected)

    self.assertArraysEqual(expected, vmapped)

  def test_vmap_fill_value_unbatched(self):
    a = self.rng().randint(0, 2, size=(3, 4, 5))
    fill_val = 99
    size = 20

    vmapped = jax.vmap(lambda x: hijax.nonzero(x, size=size, fill_value=fill_val, axes=(1,))[0])(a)

    expected = []
    for sub_a in a:
      expected.append(nonzero_reference(sub_a, size=size, fill_value=fill_val, axes=(1,))[0])
    expected = np.array(expected)

    self.assertArraysEqual(expected, vmapped)

  def test_jit_fill_value(self):
    a = self.rng().randint(0, 2, size=(100,))
    size = 150
    fill_val = 99

    @jax.jit
    def f(x, fv):
      return hijax.nonzero(x, size=size, fill_value=fv)

    actual = f(a, fill_val)
    expected = nonzero_reference(a, size=size, fill_value=fill_val)

    for e, a_i in zip(expected, actual):
      self.assertArraysEqual(e, a_i)

  def test_invalid_fill_value_shape(self):
    a = self.rng().randint(0, 2, size=(10, 20))
    size = 5
    a_aval = jax.core.ShapedArray(a.shape, a.dtype)
    fv_aval = jax.core.ShapedArray((size,), jax.numpy.int32)
    with self.assertRaisesRegex(
        ValueError, "fill_value shape .* is not broadcast-compatible with batch shape"
    ):
      hijax.Nonzero(
          a_aval,
          fv_aval,
          size=size,
          axes=(1,),
          out_dtype=np.dtype("int32"),
      )

  def test_fill_value_shape_expands_batch_shape(self):
    a = self.rng().randint(0, 2, size=(1, 20))
    size = 5
    a_aval = jax.core.ShapedArray(a.shape, a.dtype)
    fv_aval = jax.core.ShapedArray((10,), jax.numpy.int32)
    with self.assertRaisesRegex(
        ValueError, "fill_value shape .* cannot be broadcast to batch shape .* without expanding it"
    ):
      hijax.Nonzero(
          a_aval,
          fv_aval,
          size=size,
          axes=(1,),
          out_dtype=np.dtype("int32"),
      )

  @jtu.with_explicit_mesh((2, 2), ("x", "y"))
  def test_explicit_sharding(self, mesh):
    del mesh
    a_np = self.rng().randint(0, 2, size=(4, 6, 8))
    size = 5

    with self.subTest("sharded_batch_dims"):
      a = jax.device_put(a_np, P("x", "y", None))
      out = jax.jit(lambda x: hijax.nonzero(x, size=size, axes=(2,)))(a)
      self.assertEqual(jax.typeof(out[0]).sharding.spec, P("x", "y", None))
      expected = nonzero_reference(a_np, size=size, axes=(2,))
      self.assertArraysEqual(expected[0], out[0])

    with self.subTest("sharded_fill_value"):
      a_2d_np = self.rng().randint(0, 2, size=(4, 8))
      fv_np = np.array([99, 98, 97, 96], dtype=np.int32)
      a_2d = jax.device_put(a_2d_np, P("x", None))
      fv = jax.device_put(fv_np, P("x"))
      out = jax.jit(
          jax.vmap(
              lambda x, f: hijax.nonzero(x, size=10, fill_value=f, axes=(0,)),
              in_axes=(0, 0),
          )
      )(a_2d, fv)
      self.assertEqual(jax.typeof(out[0]).sharding.spec, P("x", None))
      expected = np.stack([
          nonzero_reference(sub_a, size=10, fill_value=int(f), axes=(0,))[0]
          for sub_a, f in zip(a_2d_np, fv_np)
      ])
      self.assertArraysEqual(expected, out[0])

    with self.subTest("empty_size"):
      a = jax.device_put(a_np, P("x", "y", None))
      out = jax.jit(lambda x: hijax.nonzero(x, size=0, axes=(2,)))(a)
      self.assertEqual(jax.typeof(out[0]).sharding.spec, P("x", "y", None))
      self.assertEqual(out[0].shape, (4, 6, 0))

  @jtu.with_explicit_mesh((2, 2), ("x", "y"))
  def test_explicit_sharding_errors(self, mesh):
    del mesh
    a = jax.device_put(jnp.zeros((4, 8)), P(None, "x"))
    with self.assertRaisesRegex(
        ValueError, "Cannot apply nonzero over sharded axes"
    ):
      hijax.nonzero(a, size=5, axes=(1,))

    a_1d = jax.device_put(jnp.zeros((8,)), P(None))
    fv = jax.device_put(jnp.zeros((4,), dtype=jnp.int32), P("x"))
    with self.assertRaisesRegex(
        ValueError, "fill_value sharding .* is not compatible with batch sharding"
    ):
      jax.vmap(
          lambda x, f: hijax.nonzero(x, size=5, fill_value=f, axes=(0,)),
          in_axes=(None, 0),
      )(a_1d, fv)

  @jtu.with_explicit_mesh((2, 2), ("x", "y"))
  def test_shard_map(self, mesh):
    a_np = self.rng().randint(0, 2, size=(4, 8))
    a = jax.device_put(a_np, P("x", None))

    @functools.partial(
        jax.shard_map, mesh=mesh, in_specs=(P("x", None), P()), out_specs=P("x", None)
    )
    def f(x, fv):
      return hijax.nonzero(x, size=5, fill_value=fv, axes=(1,))[0]

    out = f(a, jnp.int32(99))
    expected = nonzero_reference(a_np, size=5, fill_value=99, axes=(1,))[0]
    self.assertArraysEqual(expected, out)

class EinsumTest(jtu.JaxTestCase):

  def test_basic_against_numpy(self):
    x = self.rng().uniform(size=(2, 3))
    y = self.rng().uniform(size=(3, 4))

    expected = jnp.einsum("ij,jk->ik", x, y)
    actual = hijax.einsum("ij,jk->ik", x, y)
    self.assertAllClose(expected, actual, atol=1e-5, rtol=1e-5)

  def test_jit(self):
    x = self.rng().uniform(size=(2, 3))
    y = self.rng().uniform(size=(3, 4))

    @jax.jit
    def f(a, b):
      return hijax.einsum("ij,jk->ik", a, b)

    expected = jnp.einsum("ij,jk->ik", x, y)
    actual = f(x, y)
    self.assertAllClose(expected, actual, atol=1e-5, rtol=1e-5)

  def test_vmap(self):
    # Batching on first argument, axis 0
    x = self.rng().uniform(size=(5, 2, 3))
    y = self.rng().uniform(size=(3, 4))

    @jax.vmap
    def f(a):
      return hijax.einsum("ij,jk->ik", a, y)

    expected = jnp.stack([jnp.einsum("ij,jk->ik", a, y) for a in x])
    actual = f(x)
    self.assertAllClose(expected, actual, atol=1e-5, rtol=1e-5)

    # Batching on both
    x = self.rng().uniform(size=(5, 2, 3))
    y = self.rng().uniform(size=(5, 3, 4))

    @jax.vmap
    def g(a, b):
      return hijax.einsum("ij,jk->ik", a, b)

    expected = jnp.stack([jnp.einsum("ij,jk->ik", a, b) for a, b in zip(x, y)])
    actual = g(x, y)
    self.assertAllClose(expected, actual, atol=1e-5, rtol=1e-5)

  def test_autodiff(self):
    x = jnp.array([[1.0, 2.0], [3.0, 4.0]])
    y = jnp.array([[5.0, 6.0], [7.0, 8.0]])

    def f(a, b):
      return hijax.einsum("ij,jk->ik", a, b)

    # JVP
    tx = jnp.ones_like(x)
    ty = jnp.ones_like(y)
    primal_out, tangent_out = jax.jvp(f, (x, y), (tx, ty))

    # Expected JVP: d(X@Y) = dX @ Y + X @ dY
    expected_tangent = tx @ y + x @ ty
    self.assertAllClose(expected_tangent, tangent_out, atol=1e-5, rtol=1e-5)

    # VJP
    primal_out, f_vjp = jax.vjp(f, x, y)
    g = jnp.ones_like(primal_out)
    grads = f_vjp(g)

    # Expected VJP:
    # dL/dX = g @ Y^T
    # dL/dY = X^T @ g
    expected_grad_x = g @ y.T
    expected_grad_y = x.T @ g

    self.assertAllClose(expected_grad_x, grads[0], atol=1e-5, rtol=1e-5)
    self.assertAllClose(expected_grad_y, grads[1], atol=1e-5, rtol=1e-5)

  def test_autodiff_contracted_dim(self):
    x = jnp.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])

    def f_hijax(a):
      return hijax.einsum("ij->i", a)

    def f_jnp(a):
      return jnp.einsum("ij->i", a)

    _, vjp_hijax = jax.vjp(f_hijax, x)
    _, vjp_jnp = jax.vjp(f_jnp, x)

    g = jnp.array([1.0, 2.0])
    (grad_hijax,) = vjp_hijax(g)
    (grad_jnp,) = vjp_jnp(g)

    self.assertAllClose(grad_jnp, grad_hijax, atol=1e-5, rtol=1e-5)

    def f_scalar_hijax(a):
      return hijax.einsum("ij->", a)

    def f_scalar_jnp(a):
      return jnp.einsum("ij->", a)

    _, vjp_scalar_hijax = jax.vjp(f_scalar_hijax, x)
    _, vjp_scalar_jnp = jax.vjp(f_scalar_jnp, x)

    (grad_scalar_hijax,) = vjp_scalar_hijax(jnp.array(1.0))
    (grad_scalar_jnp,) = vjp_scalar_jnp(jnp.array(1.0))

    self.assertAllClose(grad_scalar_jnp, grad_scalar_hijax, atol=1e-5, rtol=1e-5)


  @jtu.sample_product(
      optimize=['optimal', 'greedy', False, True, 'auto'],
  )
  def test_optimize(self, optimize):
    x = self.rng().uniform(size=(2, 3))
    y = self.rng().uniform(size=(3, 4))
    z = self.rng().uniform(size=(4, 5))

    # Multiple operands to test path optimization
    expected = jnp.einsum("ij,jk,kl->il", x, y, z, optimize=optimize)
    actual = hijax.einsum("ij,jk,kl->il", x, y, z, optimize=optimize)
    self.assertAllClose(expected, actual, atol=1e-5, rtol=1e-5)

  def test_ellipsis(self):
    # Test that Ellipsis works (via parse_einsum_input normalization)
    x = self.rng().uniform(size=(2, 3, 4))
    y = self.rng().uniform(size=(4, 5))

    expected = jnp.einsum("...jk,kl->...jl", x, y)
    actual = hijax.einsum("...jk,kl->...jl", x, y)
    self.assertAllClose(expected, actual, atol=1e-5, rtol=1e-5)

  # Set the highest precision here so that on platforms like ROCm we avoid one
  # dot compiling to xf32 and the other to FMA, which makes the results disagree.
  @jax.default_matmul_precision("float32")
  def test_vmap_ellipsis(self):
    x = self.rng().uniform(size=(5, 2, 3, 4))
    y = self.rng().uniform(size=(4, 5))

    def f_raw(a, b):
      return hijax.einsum("...jk,kl->...jl", a, b)
    f = jax.vmap(f_raw, in_axes=(0, None))

    def expected_fn(a, b):
      return jnp.einsum("...jk,kl->...jl", a, b)
    expected_fn_vmap = jax.vmap(expected_fn, in_axes=(0, None))
    expected = expected_fn_vmap(x, y)

    actual = f(x, y)
    self.assertAllClose(expected, actual, atol=1e-5, rtol=1e-5)

    # Batching on non-zero axis
    x2 = self.rng().uniform(size=(2, 5, 3, 4))
    def g_raw(a, b):
      return hijax.einsum("...jk,kl->...jl", a, b)
    g = jax.vmap(g_raw, in_axes=(1, None))

    g_expected = jax.vmap(expected_fn, in_axes=(1, None))
    expected2 = g_expected(x2, y)

    actual2 = g(x2, y)
    self.assertAllClose(expected2, actual2, atol=1e-5, rtol=1e-5)

  def test_linear_transpose(self):
    x = self.rng().uniform(size=(2, 3))
    y = self.rng().uniform(size=(3, 4))

    def f(a):
      return hijax.einsum("ij,jk->ik", a, y)

    def f_ref(a):
      return jnp.einsum("ij,jk->ik", a, y)

    cotangent = jnp.ones((2, 4))

    transposed_ref_fn = jax.linear_transpose(f_ref, x)
    expected, = transposed_ref_fn(cotangent)

    transposed_fn = jax.linear_transpose(f, x)
    actual, = transposed_fn(cotangent)

    self.assertAllClose(expected, actual, atol=1e-5, rtol=1e-5)

  def test_precision(self):
    x = self.rng().uniform(size=(2, 3))
    y = self.rng().uniform(size=(3, 4))

    for precision in [None, jax.lax.Precision.DEFAULT, jax.lax.Precision.HIGH, jax.lax.Precision.HIGHEST]:
      expected = jnp.einsum("ij,jk->ik", x, y, precision=precision)
      actual = hijax.einsum("ij,jk->ik", x, y, precision=precision)
      self.assertAllClose(expected, actual, atol=1e-5, rtol=1e-5)

  def test_preferred_element_type(self):
    x = self.rng().uniform(size=(2, 3)).astype(jnp.bfloat16)
    y = self.rng().uniform(size=(3, 4)).astype(jnp.bfloat16)

    # Preferred type is float32
    expected = jnp.einsum("ij,jk->ik", x, y, preferred_element_type=np.float32)
    actual = hijax.einsum("ij,jk->ik", x, y, preferred_element_type=np.float32)
    self.assertEqual(actual.dtype, np.float32)
    self.assertAllClose(expected, actual, atol=1e-3, rtol=1e-3)

    # Preferred type is int32
    x_f32 = x.astype(np.float32)
    y_f32 = y.astype(np.float32)
    expected_int = jnp.einsum("ij,jk->ik", x_f32, y_f32, preferred_element_type=np.int32)
    actual_int = hijax.einsum("ij,jk->ik", x_f32, y_f32, preferred_element_type=np.int32)
    self.assertEqual(actual_int.dtype, np.int32)
    self.assertAllClose(expected_int, actual_int, atol=1e-5, rtol=1e-5)

  def test_vmap_none(self):
    x = self.rng().uniform(size=(2, 3))
    y = self.rng().uniform(size=(3, 4))
    f = jax.vmap(lambda a, b: hijax.einsum("ij,jk->ik", a, b), in_axes=(None, None), axis_size=5)
    actual = f(x, y)
    expected = jnp.stack([jnp.einsum("ij,jk->ik", x, y)] * 5)
    self.assertAllClose(expected, actual, atol=1e-5, rtol=1e-5)

  def test_jacobian(self):
    x = self.rng().uniform(size=(3, 4))
    y = self.rng().uniform(size=(4, 5))
    def f(x, y):
      return hijax.einsum("ij,jk->ik", x, y)
    def f_ref(x, y):
      return jnp.einsum("ij,jk->ik", x, y)

    actual_jac = jax.jacobian(f)(x, y)
    expected_jac = jax.jacobian(f_ref)(x, y)
    self.assertAllClose(expected_jac, actual_jac, atol=1e-5, rtol=1e-5)


if __name__ == "__main__":
  absltest.main(testLoader=jtu.JaxTestLoader())