Repository: dask/dask

Issue:
dask.array.ravel_multi_index doesn't replicate numpy API
**What happened**:
`dask.array.ravel_multi_index` fails when passed a tuple of integer arrays (`AttributeError: 'tuple' object has no attribute 'map_blocks'`)

**What you expected to happen**:
The documentation (copied from the equivalent numpy function) indicates that the first input should be tuple of integer arrays, one array for each dimension.

**Minimal Complete Verifiable Example**:
```python
import numpy as np
import dask.array as dsa

arr = np.array([[3,6,6],[4,5,1]]) # Example from numpy docs
darr = dsa.from_array(arr)

np.ravel_multi_index((arr, arr), (7,7))    # This works
dsa.ravel_multi_index((darr, darr), (7,7)) # This doesn't
```

**Environment**:

- Dask version: 2021.04.0
- Python version: 3.8.8
- Operating System: Linux 4.12.14-95.57-default
- Install method (conda, pip, source): conda

dask.array.ravel_multi_index doesn't replicate numpy API
**What happened**:
`dask.array.ravel_multi_index` fails when passed a tuple of integer arrays (`AttributeError: 'tuple' object has no attribute 'map_blocks'`)

**What you expected to happen**:
The documentation (copied from the equivalent numpy function) indicates that the first input should be tuple of integer arrays, one array for each dimension.

**Minimal Complete Verifiable Example**:
```python
import numpy as np
import dask.array as dsa

arr = np.array([[3,6,6],[4,5,1]]) # Example from numpy docs
darr = dsa.from_array(arr)

np.ravel_multi_index((arr, arr), (7,7))    # This works
dsa.ravel_multi_index((darr, darr), (7,7)) # This doesn't
```

**Environment**:

- Dask version: 2021.04.0
- Python version: 3.8.8
- Operating System: Linux 4.12.14-95.57-default
- Install method (conda, pip, source): conda



Discussion:
Thanks for reporting @dougiesquire! Supporting a tuple of Dask arrays as input definitely seems in scope. cc'ing @Illviljan who added the `ravel_multi_index` implementation
I've got a potential fix over at #7584. Feel free to try it out.
Thanks for reporting @dougiesquire! Supporting a tuple of Dask arrays as input definitely seems in scope. cc'ing @Illviljan who added the `ravel_multi_index` implementation
I've got a potential fix over at #7584. Feel free to try it out.

Resolving patch:
diff --git a/dask/array/routines.py b/dask/array/routines.py
--- a/dask/array/routines.py
+++ b/dask/array/routines.py
@@ -21,6 +21,7 @@
     asanyarray,
     asarray,
     blockwise,
+    broadcast_arrays,
     broadcast_shapes,
     broadcast_to,
     concatenate,
@@ -1808,18 +1809,35 @@ def unravel_index(indices, shape, order="C"):
     return unraveled_indices
 
 
-def _ravel_multi_index_kernel(multi_index, func_kwargs):
-    return np.ravel_multi_index(multi_index, **func_kwargs)
-
-
 @wraps(np.ravel_multi_index)
 def ravel_multi_index(multi_index, dims, mode="raise", order="C"):
-    return multi_index.map_blocks(
-        _ravel_multi_index_kernel,
+    if np.isscalar(dims):
+        dims = (dims,)
+    if is_dask_collection(dims) or any(is_dask_collection(d) for d in dims):
+        raise NotImplementedError(
+            f"Dask types are not supported in the `dims` argument: {dims!r}"
+        )
+
+    if is_arraylike(multi_index):
+        index_stack = asarray(multi_index)
+    else:
+        multi_index_arrs = broadcast_arrays(*multi_index)
+        index_stack = stack(multi_index_arrs)
+
+    if not np.isnan(index_stack.shape).any() and len(index_stack) != len(dims):
+        raise ValueError(
+            f"parameter multi_index must be a sequence of length {len(dims)}"
+        )
+    if not np.issubdtype(index_stack.dtype, np.signedinteger):
+        raise TypeError("only int indices permitted")
+    return index_stack.map_blocks(
+        np.ravel_multi_index,
         dtype=np.intp,
-        chunks=(multi_index.shape[-1],),
+        chunks=index_stack.chunks[1:],
         drop_axis=0,
-        func_kwargs=dict(dims=dims, mode=mode, order=order),
+        dims=dims,
+        mode=mode,
+        order=order,
     )
 
 


Test patch:
diff --git a/dask/array/tests/test_routines.py b/dask/array/tests/test_routines.py
--- a/dask/array/tests/test_routines.py
+++ b/dask/array/tests/test_routines.py
@@ -3,6 +3,8 @@
 
 import pytest
 
+from dask.delayed import delayed
+
 np = pytest.importorskip("numpy")
 
 import dask.array as da
@@ -1761,6 +1763,16 @@ def test_unravel_index():
         assert_eq(darr.vindex[d_indices], arr[indices])
 
 
+@pytest.mark.parametrize(
+    "asarray",
+    [
+        lambda x: x,
+        lambda x: [np.asarray(a) for a in x],
+        lambda x: [da.asarray(a) for a in x],
+        np.asarray,
+        da.from_array,
+    ],
+)
 @pytest.mark.parametrize(
     "arr, chunks, kwargs",
     [
@@ -1772,17 +1784,78 @@ def test_unravel_index():
         # Shape tests:
         ([[3, 6, 6]], (1, 1), dict(dims=(7), order="C")),
         ([[3, 6, 6], [4, 5, 1], [8, 6, 2]], (3, 1), dict(dims=(7, 6, 9), order="C")),
+        # Multi-dimensional index arrays
+        (
+            np.arange(6).reshape(3, 2, 1).tolist(),
+            (1, 2, 1),
+            dict(dims=(7, 6, 9), order="C"),
+        ),
+        # Broadcasting index arrays
+        ([1, [2, 3]], None, dict(dims=(8, 9))),
+        ([1, [2, 3], [[1, 2], [3, 4], [5, 6], [7, 8]]], None, dict(dims=(8, 9, 10))),
     ],
 )
-def test_ravel_multi_index(arr, chunks, kwargs):
-    arr = np.asarray(arr)
-    darr = da.from_array(arr, chunks=chunks)
+def test_ravel_multi_index(asarray, arr, chunks, kwargs):
+    if any(np.isscalar(x) for x in arr) and asarray in (np.asarray, da.from_array):
+        pytest.skip()
+
+    if asarray is da.from_array:
+        arr = np.asarray(arr)
+        input = da.from_array(arr, chunks=chunks)
+    else:
+        arr = input = asarray(arr)
+
     assert_eq(
-        da.ravel_multi_index(darr, **kwargs).compute(),
         np.ravel_multi_index(arr, **kwargs),
+        da.ravel_multi_index(input, **kwargs),
+    )
+
+
+def test_ravel_multi_index_unknown_shape():
+    multi_index = da.from_array([[3, 6, 6], [4, 5, 1], [-1, -1, -1]])
+    multi_index = multi_index[(multi_index > 0).all(axis=1)]
+
+    multi_index_np = multi_index.compute()
+
+    assert np.isnan(multi_index.shape).any()
+    assert_eq(
+        np.ravel_multi_index(multi_index_np, dims=(7, 6)),
+        da.ravel_multi_index(multi_index, dims=(7, 6)),
     )
 
 
+def test_ravel_multi_index_unknown_shape_fails():
+    multi_index1 = da.from_array([2, -1, 3, -1], chunks=2)
+    multi_index1 = multi_index1[multi_index1 > 0]
+
+    multi_index2 = da.from_array(
+        [[1, 2], [-1, -1], [3, 4], [5, 6], [7, 8], [-1, -1]], chunks=(2, 1)
+    )
+    multi_index2 = multi_index2[(multi_index2 > 0).all(axis=1)]
+
+    multi_index = [1, multi_index1, multi_index2]
+
+    assert np.isnan(multi_index1.shape).any()
+    assert np.isnan(multi_index2.shape).any()
+    with pytest.raises(ValueError, match="Arrays' chunk sizes"):
+        da.ravel_multi_index(multi_index, dims=(8, 9, 10))
+
+
+@pytest.mark.parametrize("dims", [da.from_array([5, 10]), delayed([5, 10], nout=2)])
+@pytest.mark.parametrize("wrap_in_list", [False, True])
+def test_ravel_multi_index_delayed_dims(dims, wrap_in_list):
+    with pytest.raises(NotImplementedError, match="Dask types are not supported"):
+        da.ravel_multi_index((2, 1), list(dims) if wrap_in_list else dims)
+
+
+def test_ravel_multi_index_non_int_dtype():
+    with pytest.raises(TypeError, match="only int indices permitted"):
+        da.ravel_multi_index(
+            (1.0, 2),
+            (5, 10),
+        )
+
+
 def test_coarsen():
     x = np.random.randint(10, size=(24, 24))
     d = da.from_array(x, chunks=(4, 8))