Repository: scikit-hep/awkward

Issue:
ak.* support for typetracer arrays
### Description of new feature

Running list of (discovered by trying in dask-awkward) missing functions:
- [ ] `ones_like`
- [ ] `zeros_like`
- [ ] `firsts`

Example failure:
```python
>>> import awkward._v2 as ak
>>> ak.ones_like(ak.Array(ak.Array([1, 2, 3]).layout.typetracer))
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/Users/ddavis/.pyenv/versions/dask-dev/lib/python3.10/site-packages/awkward/_v2/operations/ak_ones_like.py", line 30, in ones_like
    return _impl(array, highlevel, behavior, dtype)
  File "/Users/ddavis/.pyenv/versions/dask-dev/lib/python3.10/site-packages/awkward/_v2/operations/ak_ones_like.py", line 34, in _impl
    return ak._v2.operations.ak_full_like._impl(array, 1, highlevel, behavior, dtype)
  File "/Users/ddavis/.pyenv/versions/dask-dev/lib/python3.10/site-packages/awkward/_v2/operations/ak_full_like.py", line 177, in _impl
    out = layout.recursively_apply(action)
  File "/Users/ddavis/.pyenv/versions/dask-dev/lib/python3.10/site-packages/awkward/_v2/contents/content.py", line 1326, in recursively_apply
    return self._recursively_apply(
  File "/Users/ddavis/.pyenv/versions/dask-dev/lib/python3.10/site-packages/awkward/_v2/contents/numpyarray.py", line 1296, in _recursively_apply
    result = action(
  File "/Users/ddavis/.pyenv/versions/dask-dev/lib/python3.10/site-packages/awkward/_v2/operations/ak_full_like.py", line 164, in action
    nplike.ones_like(original),
  File "/Users/ddavis/.pyenv/versions/dask-dev/lib/python3.10/site-packages/awkward/_v2/_typetracer.py", line 531, in ones_like
    raise ak._v2._util.error(NotImplementedError)
NotImplementedError:
```


Discussion:
I made this the [top priority](https://github.com/scikit-hep/awkward/projects/5) of the "easy" column, hoping that @ioanaif, @agoose77, or @ianna can get to it soon. I think that `ones_like` and `zeros_like` will have exactly the same implementation for the TypeTracer. `firsts` is a higher-level function, but it's blocked by `ones_like` and `zeros_like`, right?

Resolving patch:
diff --git a/src/awkward/_v2/_typetracer.py b/src/awkward/_v2/_typetracer.py
--- a/src/awkward/_v2/_typetracer.py
+++ b/src/awkward/_v2/_typetracer.py
@@ -203,8 +203,22 @@ class TypeTracerArray:
     def from_array(cls, array, dtype=None):
         if isinstance(array, ak._v2.index.Index):
             array = array.data
+
+        # not array-like
+        if not hasattr(array, "shape"):
+            sequence = list(array)
+            array = numpy.array(sequence)
+            if array.dtype == np.dtype("O"):
+                raise ak._v2._util.error(
+                    ValueError(
+                        f"bug in Awkward Array: attempt to construct `TypeTracerArray` "
+                        f"from a sequence of non-primitive types: {sequence}"
+                    )
+                )
+
         if dtype is None:
             dtype = array.dtype
+
         return cls(dtype, shape=array.shape)
 
     def __init__(self, dtype, shape=None):
@@ -432,9 +446,6 @@ def copy(self):
         return self
 
 
-unset = object()
-
-
 class TypeTracer(ak.nplike.NumpyLike):
     known_data = False
     known_shape = False
@@ -483,21 +494,21 @@ def raw(self, array, nplike):
 
     ############################ array creation
 
-    def array(self, data, dtype=unset, **kwargs):
+    def array(self, data, dtype=None, **kwargs):
         # data[, dtype=[, copy=]]
-        if dtype is unset:
+        if dtype is None:
             dtype = data.dtype
         return TypeTracerArray.from_array(data, dtype=dtype)
 
-    def asarray(self, array, dtype=unset, **kwargs):
+    def asarray(self, array, dtype=None, **kwargs):
         # array[, dtype=][, order=]
-        if dtype is unset:
+        if dtype is None:
             dtype = array.dtype
         return TypeTracerArray.from_array(array, dtype=dtype)
 
-    def ascontiguousarray(self, array, dtype=unset, **kwargs):
+    def ascontiguousarray(self, array, dtype=None, **kwargs):
         # array[, dtype=]
-        if dtype is unset:
+        if dtype is None:
             dtype = array.dtype
         return TypeTracerArray.from_array(array, dtype=dtype)
 
@@ -520,23 +531,26 @@ def empty(self, shape, dtype=np.float64, **kwargs):
         # shape/len[, dtype=]
         return TypeTracerArray(dtype, shape)
 
-    def full(self, shape, value, dtype=unset, **kwargs):
+    def full(self, shape, value, dtype=None, **kwargs):
         # shape/len, value[, dtype=]
-        if dtype is unset:
+        if dtype is None:
             dtype = numpy.array(value).dtype
         return TypeTracerArray(dtype, shape)
 
-    def zeros_like(self, *args, **kwargs):
-        # array
-        raise ak._v2._util.error(NotImplementedError)
+    def zeros_like(self, a, dtype=None, **kwargs):
+        if dtype is None:
+            dtype = a.dtype
 
-    def ones_like(self, *args, **kwargs):
-        # array
-        raise ak._v2._util.error(NotImplementedError)
+        if isinstance(a, UnknownScalar):
+            return UnknownScalar(dtype)
 
-    def full_like(self, *args, **kwargs):
-        # array, fill_value
-        raise ak._v2._util.error(NotImplementedError)
+        return TypeTracerArray(dtype, a.shape)
+
+    def ones_like(self, a, dtype=None, **kwargs):
+        return self.zeros_like(a, dtype)
+
+    def full_like(self, a, fill_value, dtype=None, **kwargs):
+        return self.zeros_like(a, dtype)
 
     def arange(self, *args, **kwargs):
         # stop[, dtype=]
@@ -854,7 +868,7 @@ def argmax(self, *args, **kwargs):
         raise ak._v2._util.error(NotImplementedError)
 
     def array_str(
-        self, array, max_line_width=unset, precision=unset, suppress_small=unset
+        self, array, max_line_width=None, precision=None, suppress_small=None
     ):
         # array, max_line_width, precision=None, suppress_small=None
         return "[?? ... ??]"


Test patch:
diff --git a/tests/v2/test_1504-typetracer-like.py b/tests/v2/test_1504-typetracer-like.py
new file mode 100644
--- /dev/null
+++ b/tests/v2/test_1504-typetracer-like.py
@@ -0,0 +1,48 @@
+import awkward as ak
+import numpy as np
+import pytest
+
+
+typetracer = ak._v2._typetracer.TypeTracer.instance()
+
+
+@pytest.mark.parametrize("dtype", [np.float64, np.int64, np.uint8, None])
+@pytest.mark.parametrize("like_dtype", [np.float64, np.int64, np.uint8, None])
+def test_ones_like(dtype, like_dtype):
+    array = ak._v2.contents.numpyarray.NumpyArray(
+        np.array([99, 88, 77, 66, 66], dtype=dtype)
+    )
+    ones = ak._v2.ones_like(array.typetracer, dtype=like_dtype, highlevel=False)
+    assert ones.typetracer.shape == array.shape
+    assert ones.typetracer.dtype == like_dtype or array.dtype
+
+
+@pytest.mark.parametrize("dtype", [np.float64, np.int64, np.uint8, None])
+@pytest.mark.parametrize("like_dtype", [np.float64, np.int64, np.uint8, None])
+def test_zeros_like(dtype, like_dtype):
+    array = ak._v2.contents.numpyarray.NumpyArray(
+        np.array([99, 88, 77, 66, 66], dtype=dtype)
+    )
+
+    full = ak._v2.zeros_like(array.typetracer, dtype=like_dtype, highlevel=False)
+    assert full.typetracer.shape == array.shape
+    assert full.typetracer.dtype == like_dtype or array.dtype
+
+
+@pytest.mark.parametrize("dtype", [np.float64, np.int64, np.uint8, None])
+@pytest.mark.parametrize("like_dtype", [np.float64, np.int64, np.uint8, None])
+@pytest.mark.parametrize("value", [1.0, -20, np.iinfo(np.uint64).max])
+def test_full_like(dtype, like_dtype, value):
+    array = ak._v2.contents.numpyarray.NumpyArray(
+        np.array([99, 88, 77, 66, 66], dtype=dtype)
+    )
+    full = ak._v2.full_like(array.typetracer, value, dtype=like_dtype, highlevel=False)
+    assert full.typetracer.shape == array.shape
+    assert full.typetracer.dtype == like_dtype or array.dtype
+
+
+def test_full_like_cast():
+    with pytest.raises(ValueError):
+        ak._v2._typetracer.TypeTracerArray.from_array(
+            [1, ak._v2._typetracer.UnknownScalar(np.uint8)]
+        )