Merge branch 'main' into agoose77/fix-from-iter-non-iterable

README.md

@@ -23,6 +23,8 @@ Arrays are **dynamically typed**, but operations on them are **compiled and fast
 Given an array of lists of objects with `x`, `y` fields (with nested lists in the `y` field),
 
 ```python
+import awkward as ak
+
 array = ak.Array([
     [{"x": 1.1, "y": [1]}, {"x": 2.2, "y": [1, 2]}, {"x": 3.3, "y": [1, 2, 3]}],
     [],

src/awkward/_connect/numpy.py

@@ -5,6 +5,7 @@
 import numpy
 
 import awkward as ak
+from awkward._backends import backend_of
 from awkward._behavior import (
     behavior_of,
     find_custom_cast,
@@ -109,7 +110,7 @@ def ensure_valid_args(*args, **kwargs):
     return decorator
 
 
-def _array_ufunc_custom_cast(inputs, behavior):
+def _array_ufunc_custom_cast(inputs, behavior, backend):
     args = [
         wrap_layout(x, behavior)
         if isinstance(x, (ak.contents.Content, ak.record.Record))
@@ -122,9 +123,11 @@ def _array_ufunc_custom_cast(inputs, behavior):
         cast_fcn = find_custom_cast(x, behavior)
         if cast_fcn is not None:
             x = cast_fcn(x)
-        nextinputs.append(
-            ak.operations.to_layout(x, allow_record=True, allow_other=True)
-        )
+        maybe_layout = ak.operations.to_layout(x, allow_record=True, allow_other=True)
+        if isinstance(maybe_layout, (ak.contents.Content, ak.record.Record)):
+            maybe_layout = maybe_layout.to_backend(backend)
+
+        nextinputs.append(maybe_layout)
     return nextinputs
 
 
@@ -184,8 +187,9 @@ def array_ufunc(ufunc, method, inputs, kwargs):
         return NotImplemented
 
     behavior = behavior_of(*inputs)
+    backend = backend_of(*inputs)
 
-    inputs = _array_ufunc_custom_cast(inputs, behavior)
+    inputs = _array_ufunc_custom_cast(inputs, behavior, backend)
 
     def action(inputs, **ignore):
         signature = _array_ufunc_signature(ufunc, inputs)
@@ -205,7 +209,6 @@ def action(inputs, **ignore):
             isinstance(x, NumpyArray) or not isinstance(x, ak.contents.Content)
             for x in inputs
         ):
-            backend = ak._backends.backend_of(*inputs)
             nplike = backend.nplike
 
             # Broadcast parameters against one another

tests/test_2354_ufunc_same_backend.py

@@ -0,0 +1,23 @@
+# BSD 3-Clause License; see https://github.com/scikit-hep/awkward-1.0/blob/main/LICENSE
+
+import numpy as np  # noqa: F401
+import pytest
+
+import awkward as ak
+
+
+def test():
+    array = ak.Array([[0, 1, 2, 3], [8, 9, 10, 11]], backend="typetracer")
+    other = ak.Array([1, 2], backend="cpu")
+    result = array + other
+    assert ak.backend(result) == "typetracer"
+
+
+def test_mixed():
+    pytest.importorskip("jax")
+    ak.jax.register_and_check()
+
+    array = ak.Array([[0, 1, 2, 3], [8, 9, 10, 11]], backend="cpu")
+    other = ak.Array([1, 2], backend="jax")
+    with pytest.raises(ValueError):
+        array + other