Repository: scikit-hep/awkward

Issue:
`to_rectilinear` should support scalars
### Version of Awkward Array

1.8.0rc6

### Description and code to reproduce
```pycon
>>> x = y = ak.Array([0,1])
>>> np.ravel_multi_index((x, y), (2, 2))
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
/tmp/ipykernel_373001/4032150548.py in <module>
      1 x = y = ak.Array([0,1])
      2 
----> 3 np.ravel_multi_index((x, y), (2, 2))

<__array_function__ internals> in ravel_multi_index(*args, **kwargs)

~/.mambaforge/envs/texat/lib/python3.9/site-packages/awkward/highlevel.py in __array_function__(self, func, types, args, kwargs)
   1432         See also #__array_ufunc__.
   1433         """
-> 1434         return ak._connect._numpy.array_function(func, types, args, kwargs)
   1435 
   1436     @property

~/.mambaforge/envs/texat/lib/python3.9/site-packages/awkward/_connect/_numpy.py in array_function(func, types, args, kwargs)
     32     function = implemented.get(func)
     33     if function is None:
---> 34         args = tuple(_to_rectilinear(x) for x in args)
     35         kwargs = {k: _to_rectilinear(v) for k, v in kwargs.items()}
     36         out = func(*args, **kwargs)

~/.mambaforge/envs/texat/lib/python3.9/site-packages/awkward/_connect/_numpy.py in <genexpr>(.0)
     32     function = implemented.get(func)
     33     if function is None:
---> 34         args = tuple(_to_rectilinear(x) for x in args)
     35         kwargs = {k: _to_rectilinear(v) for k, v in kwargs.items()}
     36         out = func(*args, **kwargs)

~/.mambaforge/envs/texat/lib/python3.9/site-packages/awkward/_connect/_numpy.py in _to_rectilinear(arg)
     24     if isinstance(arg, Iterable):
     25         nplike = ak.nplike.of(arg)
---> 26         return nplike.to_rectilinear(arg, allow_missing=False)
     27     else:
     28         return arg

~/.mambaforge/envs/texat/lib/python3.9/site-packages/awkward/nplike.py in to_rectilinear(self, array, *args, **kwargs)
    447 
    448         elif isinstance(array, Iterable):
--> 449             return [self.to_rectilinear(x, *args, **kwargs) for x in array]
    450 
    451         else:

~/.mambaforge/envs/texat/lib/python3.9/site-packages/awkward/nplike.py in <listcomp>(.0)
    447 
    448         elif isinstance(array, Iterable):
--> 449             return [self.to_rectilinear(x, *args, **kwargs) for x in array]
    450 
    451         else:

~/.mambaforge/envs/texat/lib/python3.9/site-packages/awkward/nplike.py in to_rectilinear(self, array, *args, **kwargs)
    450 
    451         else:
--> 452             raise TypeError("to_rectilinear argument must be iterable")
    453 
    454     def __getitem__(self, name_and_types):

TypeError: to_rectilinear argument must be iterable
```

I don't think this is related to a6e04dd147e2993944c56255928d6f0e2e25fd49 although I didn't see this error before. I *think* this might be related to #1198.

There are some NumPy functions that receive iterable arguments of scalars, e.g `np.ravel_multi_index`. It is possible to work-around the bug in this case by passing in an `Array` instead of tuple-of-scalars, but I think that we should support this.



Discussion:

Resolving patch:
diff --git a/src/awkward/_connect/numpy.py b/src/awkward/_connect/numpy.py
--- a/src/awkward/_connect/numpy.py
+++ b/src/awkward/_connect/numpy.py
@@ -25,12 +25,23 @@ def convert_to_array(layout, args, kwargs):
 
 
 def _to_rectilinear(arg):
-    if isinstance(arg, tuple):
-        nplike = ak._nplikes.nplike_of(*arg)
-        return tuple(nplike.to_rectilinear(x) for x in arg)
+    backend = ak._backends.backend_of(arg, default=None)
+    # We have some array-like object that our backend mechanism understands
+    if backend is not None:
+        return backend.nplike.to_rectilinear(arg)
+    elif isinstance(arg, tuple):
+        return tuple(_to_rectilinear(x) for x in arg)
+    elif isinstance(arg, list):
+        return [_to_rectilinear(x) for x in arg]
+    elif ak._util.is_non_string_iterable(arg):
+        raise ak._errors.wrap_error(
+            TypeError(
+                f"encountered an unsupported iterable value {arg!r} whilst converting arguments to NumPy-friendly "
+                f"types. If this argument should be supported, please file a bug report."
+            )
+        )
     else:
-        nplike = ak._nplikes.nplike_of(arg)
-        return nplike.to_rectilinear(arg)
+        return arg
 
 
 def array_function(func, types, args, kwargs, behavior):
diff --git a/src/awkward/_nplikes.py b/src/awkward/_nplikes.py
--- a/src/awkward/_nplikes.py
+++ b/src/awkward/_nplikes.py
@@ -2,10 +2,8 @@
 from __future__ import annotations
 
 import ctypes
-from collections.abc import Iterable
 
 import numpy
-from awkward_cpp.lib import _ext
 
 import awkward as ak
 from awkward.typing import TypeVar
@@ -382,6 +380,9 @@ def is_own_array(cls, obj) -> bool:
     def is_c_contiguous(self, array) -> bool:
         raise ak._errors.wrap_error(NotImplementedError)
 
+    def to_rectilinear(self, array):
+        raise ak._errors.wrap_error(NotImplementedError)
+
 
 class Kernel:
     def __init__(self, kernel, name_and_types):
@@ -482,29 +483,7 @@ def __call__(self, *args):
 
 class Numpy(NumpyLike):
     def to_rectilinear(self, array, *args, **kwargs):
-        if isinstance(array, numpy.ndarray):
-            return array
-
-        elif isinstance(
-            array,
-            (
-                ak.Array,
-                ak.Record,
-                ak.ArrayBuilder,
-                ak.contents.Content,
-                ak.record.Record,
-                _ext.ArrayBuilder,
-            ),
-        ):
-            return ak.operations.ak_to_numpy.to_numpy(array, *args, **kwargs)
-
-        elif isinstance(array, Iterable):
-            return [self.to_rectilinear(x, *args, **kwargs) for x in array]
-
-        else:
-            raise ak._errors.wrap_error(
-                TypeError("to_rectilinear argument must be iterable")
-            )
+        return ak.operations.ak_to_numpy.to_numpy(array, *args, **kwargs)
 
     def __init__(self):
         self._module = numpy
@@ -761,29 +740,7 @@ def is_c_contiguous(self, array) -> bool:
 
 class Jax(NumpyLike):
     def to_rectilinear(self, array, *args, **kwargs):
-        if isinstance(array, self._module.DeviceArray):
-            return array
-
-        elif isinstance(
-            array,
-            (
-                ak.Array,
-                ak.Record,
-                ak.ArrayBuilder,
-                ak.contents.Content,
-                ak.record.Record,
-                _ext.ArrayBuilder,
-            ),
-        ):
-            return ak.operations.ak_to_jax.to_jax(array, *args, **kwargs)
-
-        elif isinstance(array, Iterable):
-            return [self.to_rectilinear(x, *args, **kwargs) for x in array]
-
-        else:
-            raise ak._errors.wrap_error(
-                ValueError("to_rectilinear argument must be iterable")
-            )
+        return ak.operations.ak_to_jax.to_jax(array, *args, **kwargs)
 
     def __init__(self):
         jax = ak.jax.import_jax()


Test patch:
diff --git a/tests/test_1318_array_function_types.py b/tests/test_1318_array_function_types.py
new file mode 100644
--- /dev/null
+++ b/tests/test_1318_array_function_types.py
@@ -0,0 +1,44 @@
+# BSD 3-Clause License; see https://github.com/scikit-hep/awkward-1.0/blob/main/LICENSE
+import numpy as np
+
+import awkward as ak
+
+
+def test_tuple():
+    data = (np.array([0, 1], dtype=np.int64), np.array([1, 0], dtype=np.int64))
+    result = np.ravel_multi_index(
+        (ak.from_numpy(data[0]), ak.from_numpy(data[1])), (2, 2)
+    )
+    assert isinstance(result, ak.Array)
+    assert ak._util.arrays_approx_equal(result, np.ravel_multi_index(data, (2, 2)))
+
+
+def test_list():
+    data = np.array([1, 2, 3, 4, 3, 2, 1, 2], dtype=np.int64)
+    result = np.partition(ak.from_numpy(data), [4, 6])
+    assert isinstance(result, ak.Array)
+    assert ak._util.arrays_approx_equal(result, np.partition(data, [4, 6]))
+
+
+def test_array():
+    data = np.array([1, 2, 3, 4, 3, 2, 1, 2], dtype=np.int64)
+    result = np.partition(ak.from_numpy(data), ak.Array([4, 6]))
+    assert isinstance(result, ak.Array)
+    assert ak._util.arrays_approx_equal(result, np.partition(data, np.array([4, 6])))
+
+
+def test_scalar():
+    data = np.array([1, 2, 3, 4, 3, 2, 1, 2], dtype=np.int64)
+    result = np.partition(ak.from_numpy(data), 4)
+    assert isinstance(result, ak.Array)
+    assert ak._util.arrays_approx_equal(result, np.partition(data, 4))
+
+
+def test_tuple_of_array():
+    data = (
+        np.array([5, 4, 3, 2, 1], dtype=np.int64),
+        np.array([1, 2, 3, 5, 0], dtype=np.int64),
+    )
+    result = np.lexsort((ak.from_numpy(data[0]), ak.from_numpy(data[1])))
+    assert isinstance(result, ak.Array)
+    assert ak._util.arrays_approx_equal(result, np.lexsort(data))