refactor: move typetracer ufunc handling to backend [1 of 2]

src/awkward/_backends.py

@@ -12,6 +12,7 @@
 from awkward.typing import Callable, Final, Tuple, TypeAlias, TypeVar, Unpack
 
 np = NumpyMetadata.instance()
+numpy = ak._nplikes.Numpy.instance()
 
 
 T = TypeVar("T", covariant=True)
@@ -20,7 +21,10 @@
 
 
 class Backend(Singleton, ABC):
-    name: str
+    @property
+    @abstractmethod
+    def name(self) -> str:
+        raise ak._errors.wrap_error(NotImplementedError)
 
     @property
     @abstractmethod
@@ -163,6 +167,34 @@ def __init__(self):
     def __getitem__(self, index: KernelKeyType) -> TypeTracerKernel:
         return TypeTracerKernel(index)
 
+    def _coerce_ufunc_argument(self, x):
+        if isinstance(x, ak._typetracer.TypeTracerArray):
+            return numpy.empty((0,) + x.shape[1:], dtype=x.dtype)
+        # Convert scalars to 0-d arrays
+        elif isinstance(x, ak._typetracer.UnknownScalar):
+            return numpy.empty((0,), dtype=x.dtype)
+        elif x is ak._typetracer.UnknownLength:
+            return numpy.empty((0,), dtype=np.int64)
+        elif isinstance(x, ak._typetracer.MaybeNone):
+            return self._coerce_ufunc_argument(x.content)
+        else:
+            return x
+
+    def apply_ufunc(self, ufunc, method, args, kwargs):
+        shape = None
+        numpy_args = []
+
+        for x in args:
+            if isinstance(x, ak._typetracer.TypeTracerArray):
+                x.touch_data()
+                shape = x.shape
+
+            numpy_args.append(self._coerce_ufunc_argument(x))
+
+        assert shape is not None
+        tmp = getattr(ufunc, method)(*numpy_args, **kwargs)
+        return self._typetracer.empty((shape[0],) + tmp.shape[1:], dtype=tmp.dtype)
+
 
 def _backend_for_nplike(nplike: ak._nplikes.NumpyLike) -> Backend:
     # Currently there exists a one-to-one relationship between the nplike

src/awkward/_connect/numpy.py

@@ -209,32 +209,16 @@ def action(inputs, **ignore):
                 inputs
             )
             (parameters,) = parameters_factory(1)
-            if nplike.known_data:
-                args = []
-                for x in inputs:
-                    if isinstance(x, NumpyArray):
-                        args.append(x._raw(nplike))
-                    else:
-                        args.append(x)
 
-                result = backend.apply_ufunc(ufunc, method, args, kwargs)
+            args = []
+            for x in inputs:
+                if isinstance(x, NumpyArray):
+                    args.append(x._raw(nplike))
+                else:
+                    args.append(x)
+
+            result = backend.apply_ufunc(ufunc, method, args, kwargs)
 
-            else:
-                shape = None
-                args = []
-                for x in inputs:
-                    if isinstance(x, NumpyArray):
-                        # some ufuncs have multiple array arguments, and they might
-                        # not all be typetracers
-                        if isinstance(x.data, ak._typetracer.TypeTracerArray):
-                            x.data.touch_data()
-                        shape = x.shape
-                        args.append(numpy.empty((0,) + x.shape[1:], dtype=x.dtype))
-                    else:
-                        args.append(x)
-                assert shape is not None
-                tmp = getattr(ufunc, method)(*args, **kwargs)
-                result = nplike.empty((shape[0],) + tmp.shape[1:], dtype=tmp.dtype)
             return (NumpyArray(result, backend=backend, parameters=parameters),)
 
         for x in inputs:

tests/test_2150_typetracer_high_level_ufunc.py

@@ -0,0 +1,53 @@
+# BSD 3-Clause License; see https://github.com/scikit-hep/awkward-1.0/blob/main/LICENSE
+
+import numpy as np
+import pytest  # noqa: F401
+
+import awkward as ak
+
+
+def test():
+    array = ak.Array([1, 2, 3, 4])
+    typetracer_array = ak.Array(array.layout.to_typetracer(forget_length=True))
+    typetracer_result = np.sqrt(typetracer_array)
+
+    assert typetracer_result.type == ak.types.ArrayType(
+        ak.types.NumpyType("float64"), ak._typetracer.UnknownLength
+    )
+
+
+def test_add():
+    left = ak.Array([1, 2, 3, 4])
+    right = ak.Array([1, 2, 3, 4.0])
+    typetracer_left = ak.Array(left.layout.to_typetracer(forget_length=True))
+    typetracer_right = ak.Array(right.layout.to_typetracer(forget_length=True))
+    typetracer_result = np.add(typetracer_left, typetracer_right)
+
+    assert typetracer_result.type == ak.types.ArrayType(
+        ak.types.NumpyType("float64"), ak._typetracer.UnknownLength
+    )
+
+
+def test_add_scalar():
+    array = ak.Array([1, 2, 3, 4])
+    typetracer_array = ak.Array(array.layout.to_typetracer(forget_length=True))
+    other = ak.min(typetracer_array, mask_identity=False, initial=10)
+    assert isinstance(other, ak._typetracer.UnknownScalar)
+
+    typetracer_result = np.add(typetracer_array, other)
+    assert typetracer_result.type == ak.types.ArrayType(
+        ak.types.NumpyType("int64"), ak._typetracer.UnknownLength
+    )
+
+
+def test_add_none_scalar():
+    array = ak.Array([1, 2, 3, 4])
+    typetracer_array = ak.Array(array.layout.to_typetracer(forget_length=True))
+    other = ak.min(typetracer_array, mask_identity=True, initial=10)
+    assert isinstance(other, ak._typetracer.MaybeNone)
+    assert isinstance(other.content, ak._typetracer.UnknownScalar)
+
+    typetracer_result = np.add(typetracer_array, other)
+    assert typetracer_result.type == ak.types.ArrayType(
+        ak.types.NumpyType("int64"), ak._typetracer.UnknownLength
+    )