fix core_indices

Author: crusaderky

src/array_api_extra/_apply.py

@@ -32,7 +32,7 @@ def apply_numpy_func(  # type: ignore[no-any-explicit]
     dtypes: Sequence[DType] | None = None,
     xp: ModuleType | None = None,
     input_indices: Sequence[Sequence[Hashable]] | None = None,
-    core_indices: Sequence[Sequence[Hashable]] | None = None,
+    core_indices: Sequence[Hashable] | None = None,
     output_indices: Sequence[Sequence[Hashable]] | None = None,
     adjust_chunks: Sequence[dict[Hashable, Callable[[int], int]]] | None = None,
     new_axes: Sequence[dict[Hashable, int]] | None = None,
@@ -70,9 +70,9 @@ def apply_numpy_func(  # type: ignore[no-any-explicit]
         ndim=3 and 1, `input_indices` could be ``['ijk', 'j']`` or ``[(0, 1, 2),
         (1,)]``.
         Default: disallow Dask.
-    core_indices : Sequence[Sequence[Hashable]], optional
+    core_indices : Sequence[Hashable], optional
         **Dask specific.**
-        Axes labels of each input array that cannot be broken into chunks.
+        Axes of the input arrays that cannot be broken into chunks.
         Default: disallow Dask.
     output_indices : Sequence[Sequence[Hashable]], optional
         **Dask specific.**
@@ -144,7 +144,7 @@ def apply_numpy_func(  # type: ignore[no-any-explicit]
 
             >>> apply_numpy_func(lambda x: x + x.sum(axis=0), x,
             ...                  input_indices=['ij'], output_indices=['ij'],
-            ...                  core_indices=['i'])
+            ...                  core_indices='i')
 
         This will cause `apply_numpy_func` to raise if the first axis of `x` is broken
         along multiple chunks.
@@ -177,9 +177,6 @@ def apply_numpy_func(  # type: ignore[no-any-explicit]
         if len(input_indices) != len(args):
             msg = f"got {len(input_indices)} input_indices and {len(args)} args"
             raise ValueError(msg)
-        if len(core_indices) != len(args):
-            msg = f"got {len(core_indices)} input_indices and {len(args)} args"
-            raise ValueError(msg)
         if len(output_indices) != len(shapes):
             msg = f"got {len(output_indices)} input_indices and {len(shapes)} shapes"
             raise NotImplementedError(msg)
@@ -197,19 +194,9 @@ def apply_numpy_func(  # type: ignore[no-any-explicit]
             raise ValueError(msg)
 
         # core_indices validation
-        for core_idx, inp_idx, arg in zip(
-            core_indices, input_indices, args, strict=True
-        ):
-            for i in core_idx:
-                try:
-                    axis = list(inp_idx).index(i)
-                except ValueError:
-                    msg = (
-                        f"Index {i} found in core indices but not in "
-                        "matching input_indices"
-                    )
-                    raise ValueError(msg) from None
-                if len(arg.chunks[axis]) > 1:
+        for inp_idx, arg in zip(input_indices, args, strict=True):
+            for i, chunks in zip(inp_idx, arg.chunks, strict=True):
+                if i in core_indices and len(chunks) > 1:
                     msg = f"Core index {i} is broken into multiple chunks"
                     raise ValueError(msg)