fix: backends should have defaults for user-facing operations (#1940)

src/awkward/_backends.py

@@ -21,6 +21,7 @@
 from awkward.typing import (
     Any,
     Callable,
+    Final,
     Protocol,
     Self,
     Tuple,
@@ -40,6 +41,8 @@
 
 @runtime_checkable
 class Backend(Protocol[T]):
+    name: str
+
     @property
     @abstractmethod
     def nplike(self) -> NumpyLike:
@@ -60,6 +63,8 @@ def __getitem__(self, key: KernelKeyType) -> KernelType:
 
 
 class NumpyBackend(Singleton, Backend[Any]):
+    name: Final[str] = "cpu"
+
     _numpy: Numpy
 
     @property
@@ -78,6 +83,8 @@ def __getitem__(self, index: KernelKeyType) -> NumpyKernel:
 
 
 class CupyBackend(Singleton, Backend[Any]):
+    name: Final[str] = "cuda"
+
     _cupy: Cupy
 
     @property
@@ -106,6 +113,8 @@ def __getitem__(self, index: KernelKeyType) -> CupyKernel | NumpyKernel:
 
 
 class JaxBackend(Singleton, Backend[Any]):
+    name: Final[str] = "jax"
+
     _jax: Jax
     _numpy: Numpy
 
@@ -127,6 +136,8 @@ def __getitem__(self, index: KernelKeyType) -> JaxKernel:
 
 
 class TypeTracerBackend(Singleton, Backend[Any]):
+    name: Final[str] = "typetracer"
+
     _typetracer: TypeTracer
 
     @property
@@ -182,10 +193,8 @@ def backend_of(*objects, default: D = _UNSET) -> Backend | D:
         return default
 
 
-_backends = {
-    "cpu": NumpyBackend,
-    "cuda": CupyBackend,
-    "jax": JaxBackend,
+_backends: Final[dict[str, type[Backend]]] = {
+    b.name: b for b in (NumpyBackend, CupyBackend, JaxBackend, TypeTracerBackend)
 }
 
 
@@ -195,6 +204,4 @@ def regularize_backend(backend: str | Backend) -> Backend:
     elif backend in _backends:
         return _backends[backend].instance()
     else:
-        raise ak._errors.wrap_error(
-            ValueError("The available backends for now are `cpu` and `cuda`.")
-        )
+        raise ak._errors.wrap_error(ValueError(f"No such backend {backend!r} exists."))

src/awkward/operations/ak_backend.py

@@ -3,33 +3,16 @@
 import awkward as ak
 
 
-def backend(*arrays):
+def backend(*arrays) -> str:
     """
     Returns the names of the backend used by `arrays`. May be
 
-       * `"cpu"` for `libawkward-cpu-kernels.so`;
-       * `"cuda"` for `libawkward-cuda-kernels.so`;
-       * `"mixed"` if any of the arrays have different labels within their
-         structure or any arrays have different labels from each other;
-       * None if the objects are not Awkward, NumPy, or CuPy arrays (e.g.
+       * `"cpu"` for arrays backed by NumPy;
+       * `"cuda"` for arrays backed by CuPy;
+       * `"jax"` for arrays backed by JAX;
+       * None if the objects are not Awkward, NumPy, JAX, or CuPy arrays (e.g.
          Python numbers, booleans, strings).
 
-    Mixed arrays can't be used in any operations, and two arrays on different
-    devices can't be used in the same operation.
-
-    To use `"cuda"`, the package
-    [awkward-cuda-kernels](https://pypi.org/project/awkward-cuda-kernels)
-    be installed, either by
-
-        pip install awkward-cuda-kernels
-
-    or as an optional dependency with
-
-        pip install awkward[cuda] --upgrade
-
-    It is only available for Linux as a binary wheel, and only supports Nvidia
-    GPUs (it is written in CUDA).
-
     See #ak.to_backend.
     """
     with ak._errors.OperationErrorContext(
@@ -39,32 +22,14 @@ def backend(*arrays):
         return _impl(arrays)
 
 
-def _impl(arrays):
-    backends = set()
-    for array in arrays:
-        layout = ak.operations.to_layout(
-            array,
-            allow_record=True,
-            allow_other=True,
+def _impl(arrays) -> str:
+    backend_impl = ak._backends.backend_of(*arrays, default=None)
+    if isinstance(backend_impl, ak._backends.TypeTracerBackend):
+        raise ak._errors.wrap_error(
+            ValueError(
+                "at least one of the given arrays was a typetracer array. "
+                "This is an internal backend that you should not have encountered. "
+                "Please file a bug report at https://github.com/scikit-hep/awkward/issues/"
+            )
         )
-        # Find the nplike, if it is explicitly associated with this object
-        nplike = ak.nplikes.nplike_of(layout, default=None)
-        if nplike is None:
-            continue
-        if isinstance(nplike, ak.nplikes.Jax):
-            backends.add("jax")
-        elif isinstance(nplike, ak.nplikes.Cupy):
-            backends.add("cuda")
-        elif isinstance(nplike, ak.nplikes.Numpy):
-            backends.add("cpu")
-
-    if backends == set():
-        return None
-    elif backends == {"cpu"}:
-        return "cpu"
-    elif backends == {"cuda"}:
-        return "cuda"
-    elif backends == {"jax"}:
-        return "jax"
-    else:
-        return "mixed"
+    return backend_impl.name

src/awkward/operations/ak_cartesian.py

@@ -3,6 +3,7 @@
 import awkward as ak
 
 np = ak.nplikes.NumpyMetadata.instance()
+cpu = ak._backends.NumpyBackend.instance()
 
 
 def cartesian(
@@ -233,7 +234,7 @@ def cartesian(
 def _impl(arrays, axis, nested, parameters, with_name, highlevel, behavior):
     if isinstance(arrays, dict):
         behavior = ak._util.behavior_of(*arrays.values(), behavior=behavior)
-        backend = ak._backends.backend_of(*arrays.values())
+        backend = ak._backends.backend_of(*arrays.values(), default=cpu)
         new_arrays = {}
         for n, x in arrays.items():
             new_arrays[n] = ak.operations.to_layout(
@@ -243,7 +244,7 @@ def _impl(arrays, axis, nested, parameters, with_name, highlevel, behavior):
     else:
         arrays = list(arrays)
         behavior = ak._util.behavior_of(*arrays, behavior=behavior)
-        backend = ak._backends.backend_of(*arrays)
+        backend = ak._backends.backend_of(*arrays, default=cpu)
         new_arrays = []
         for x in arrays:
             new_arrays.append(

src/awkward/operations/ak_concatenate.py

@@ -4,6 +4,7 @@
 from awkward.operations.ak_fill_none import fill_none
 
 np = ak.nplikes.NumpyMetadata.instance()
+cpu = ak._backends.NumpyBackend.instance()
 
 
 @ak._connect.numpy.implements("concatenate")
@@ -140,7 +141,7 @@ def action(inputs, depth, **kwargs):
                 inputs = nextinputs
 
             if depth == posaxis:
-                backend = ak._backends.backend_of(*inputs)
+                backend = ak._backends.backend_of(*inputs, default=cpu)
 
                 length = ak._typetracer.UnknownLength
                 for x in inputs:

src/awkward/operations/ak_fill_none.py

@@ -5,6 +5,7 @@
 import awkward as ak
 
 np = ak.nplikes.NumpyMetadata.instance()
+cpu = ak._backends.NumpyBackend.instance()
 
 
 def fill_none(array, value, axis=-1, *, highlevel=True, behavior=None):
@@ -60,7 +61,7 @@ def fill_none(array, value, axis=-1, *, highlevel=True, behavior=None):
 def _impl(array, value, axis, highlevel, behavior):
     arraylayout = ak.operations.to_layout(array, allow_record=True, allow_other=False)
     behavior = ak._util.behavior_of(array, behavior=behavior)
-    backend = ak._backends.backend_of(arraylayout)
+    backend = ak._backends.backend_of(arraylayout, default=cpu)
 
     # Convert value type to appropriate layout
     if (

src/awkward/operations/ak_run_lengths.py

@@ -3,6 +3,7 @@
 import awkward as ak
 
 np = ak.nplikes.NumpyMetadata.instance()
+cpu = ak._backends.NumpyBackend.instance()
 
 
 def run_lengths(array, *, highlevel=True, behavior=None):
@@ -95,9 +96,7 @@ def run_lengths(array, *, highlevel=True, behavior=None):
 
 
 def _impl(array, highlevel, behavior):
-    backend = ak._backends.backend_of(
-        array, default=ak._backends.NumpyBackend.instance()
-    )
+    backend = ak._backends.backend_of(array, default=cpu)
 
     def lengths_of(data, offsets):
         if len(data) == 0:

src/awkward/operations/ak_to_backend.py

@@ -9,40 +9,38 @@ def to_backend(array, backend, *, highlevel=True, behavior=None):
     """
     Args:
         array: Data to convert to a specified `backend` set.
-        backend (`"cpu"` or `"cuda"`): If `"cpu"`, the array structure is
+        backend (`"cpu"`, `"cuda"`, or `"jax"`): If `"cpu"`, the array structure is
             recursively copied (if need be) to main memory for use with
-            the default `libawkward-cpu-kernels.so`; if `"cuda"`, the
-            structure is copied to the GPU(s) for use with
-            `libawkward-cuda-kernels.so`.
+            the default Numpy backend; if `"cuda"`, the structure is copied
+            to the GPU(s) for use with CuPy. If `"jax"`, the structure is
+            copied to the CPU for use with JAX.
         highlevel (bool): If True, return an #ak.Array; otherwise, return
             a low-level #ak.contents.Content subclass.
         behavior (None or dict): Custom #ak.behavior for the output array, if
             high-level.
 
-    Converts an array from `"cpu"`, `"cuda"`, or `"mixed"` kernels to `"cpu"`
-    or `"cuda"`.
-
-    An array is `"mixed"` if some components are set to use the `"cpu"` backend and
-    others are set to use the `"cuda"` backend. Mixed arrays can't be used in any
-    operations, and two arrays set to different backends can't be used in the
-    same operation.
+    Converts an array from `"cpu"`, `"cuda"`, or `"jax"` kernels to `"cpu"`,
+    `"cuda"`, or `"jax"`.
 
     Any components that are already in the desired backend are viewed,
     rather than copied, so this operation can be an inexpensive way to ensure
     that an array is ready for a particular library.
 
-    To use `"cuda"`, the package
-    [awkward-cuda-kernels](https://pypi.org/project/awkward-cuda-kernels)
-    be installed, either by
+    To use `"cuda"`, the `cupy` package must be installed, either with
+
+        pip install cupy
+
+    or
+
+        conda install -c conda-forge cupy
 
-        pip install awkward-cuda-kernels
+    To use `"jax"`, the `jax` package must be installed, either with
 
-    or as an optional dependency with
+        pip install jax
 
-        pip install awkward[cuda] --upgrade
+    or
 
-    It is only available for Linux as a binary wheel, and only supports Nvidia
-    GPUs (it is written in CUDA).
+        conda install -c conda-forge jax
 
     See #ak.kernels.
     """

src/awkward/operations/ak_transform.py

@@ -5,6 +5,8 @@
 
 import awkward as ak
 
+cpu = ak._backends.NumpyBackend.instance()
+
 
 def transform(
     transformation,
@@ -441,7 +443,7 @@ def _impl(
     more_layouts = [
         ak.to_layout(x, allow_record=False, allow_other=False) for x in more_arrays
     ]
-    backend = ak._backends.backend_of(layout, *more_layouts)
+    backend = ak._backends.backend_of(layout, *more_layouts, default=cpu)
 
     options = {
         "allow_records": allow_records,

src/awkward/operations/ak_where.py

@@ -3,6 +3,7 @@
 import awkward as ak
 
 np = ak.nplikes.NumpyMetadata.instance()
+cpu = ak._backends.NumpyBackend.instance()
 
 
 @ak._connect.numpy.implements("where")
@@ -77,7 +78,7 @@ def _impl1(condition, mergebool, highlevel, behavior):
     akcondition = ak.operations.to_layout(
         condition, allow_record=False, allow_other=False
     )
-    backend = ak._backends.backend_of(akcondition)
+    backend = ak._backends.backend_of(akcondition, default=cpu)
 
     akcondition = ak.contents.NumpyArray(ak.operations.to_numpy(akcondition))
     out = backend.nplike.nonzero(ak.operations.to_numpy(akcondition))
@@ -106,7 +107,7 @@ def _impl3(condition, x, y, mergebool, highlevel, behavior):
         good_arrays.append(left)
     if isinstance(right, ak.contents.Content):
         good_arrays.append(right)
-    backend = ak._backends.backend_of(*good_arrays)
+    backend = ak._backends.backend_of(*good_arrays, default=cpu)
 
     def action(inputs, **kwargs):
         akcondition, left, right = inputs

tests/test_1940-ak-backend.py

@@ -0,0 +1,22 @@
+# BSD 3-Clause License; see https://github.com/scikit-hep/awkward-1.0/blob/main/LICENSE
+
+import numpy as np  # noqa: F401
+import pytest  # noqa: F401
+
+import awkward as ak
+
+
+def test_typetracer():
+    array = ak.Array([[0, 1, 2, 3], [8, 9, 10, 11]])
+    typetracer = ak.Array(array.layout.typetracer)
+
+    with pytest.raises(ValueError, match="internal backend"):
+        ak.backend(typetracer)
+
+
+def test_typetracer_mixed():
+    array = ak.Array([[0, 1, 2, 3], [8, 9, 10, 11]])
+    typetracer = ak.Array(array.layout.typetracer)
+
+    with pytest.raises(ValueError, match="internal backend"):
+        ak.backend(typetracer, array)