feat: implement virtual arrays

src/awkward/_kernels.py

@@ -12,6 +12,7 @@
 from awkward._nplikes.numpy import Numpy
 from awkward._nplikes.numpy_like import NumpyMetadata
 from awkward._nplikes.typetracer import try_touch_data
+from awkward._nplikes.virtual import materialize_if_virtual
 from awkward._typing import Protocol, TypeAlias
 
 KernelKeyType: TypeAlias = tuple  # Tuple[str, Unpack[Tuple[metadata.dtype, ...]]]
@@ -88,6 +89,8 @@ def _cast(cls, x, t):
     def __call__(self, *args) -> None:
         assert len(args) == len(self._impl.argtypes)
 
+        args = materialize_if_virtual(*args)
+
         return self._impl(
             *(self._cast(x, t) for x, t in zip(args, self._impl.argtypes))
         )
@@ -97,6 +100,8 @@ class JaxKernel(NumpyKernel):
     def __call__(self, *args) -> None:
         assert len(args) == len(self._impl.argtypes)
 
+        args = materialize_if_virtual(*args)
+
         if not any(Jax.is_tracer_type(type(arg)) for arg in args):
             return super().__call__(*args)
 
@@ -138,6 +143,8 @@ def _cast(self, x, type_):
     def __call__(self, *args) -> None:
         import awkward._connect.cuda as ak_cuda
 
+        args = materialize_if_virtual(*args)
+
         cupy = ak_cuda.import_cupy("Awkward Arrays with CUDA")
         maxlength = self.max_length(args)
         grid, blocks = self.calc_grid(maxlength), self.calc_blocks(maxlength)

src/awkward/_nplikes/__init__.py

@@ -6,6 +6,7 @@
 import awkward._nplikes.jax
 import awkward._nplikes.numpy
 import awkward._nplikes.typetracer
+import awkward._nplikes.virtual
 from awkward._nplikes.dispatch import nplike_of_obj
 from awkward._typing import TYPE_CHECKING