fix: support N-d sequences in TypeTracer.asarray (#2361)

src/awkward/_nplikes/typetracer.py

@@ -676,10 +676,52 @@ def asarray(
                 return TypeTracerArray._new(as_array.dtype, ())
 
             elif is_non_string_like_sequence(obj):
-                assert not any(is_non_string_like_sequence(x) for x in obj)
-                shape = (len(obj),)
-                result_type = numpy.result_type(*obj)  # TODO: result_type
-                return TypeTracerArray._new(result_type, shape)
+                shape = []
+                flat_items = []
+                has_seen_leaf = False
+
+                # DFS walk into sequence, construct shape, then validate
+                # remainder of the sequence against this shape.
+                def populate_shape_and_items(node, dim):
+                    nonlocal has_seen_leaf
+
+                    # If we've already computed the shape,
+                    # ensure this item matches!
+                    if has_seen_leaf:
+                        if len(node) != shape[dim - 1]:
+                            raise wrap_error(
+                                ValueError(
+                                    f"sequence at dimension {dim} does not match shape {shape[dim-1]}"
+                                )
+                            )
+                    else:
+                        shape.append(len(node))
+
+                    if isinstance(node, TypeTracerArray):
+                        raise wrap_error(
+                            AssertionError(
+                                "typetracer arrays inside sequences not currently supported"
+                            )
+                        )
+                    # Found leaf!
+                    elif len(node) == 0 or not is_non_string_like_sequence(node[0]):
+                        has_seen_leaf = True
+                        flat_items.extend(
+                            [
+                                item.dtype if is_unknown_scalar(item) else item
+                                for item in node
+                            ]
+                        )
+
+                    # Keep recursing!
+                    else:
+                        for child in node:
+                            populate_shape_and_items(child, dim + 1)
+
+                populate_shape_and_items(obj, 1)
+                if dtype is None:
+                    dtype = numpy.result_type(*flat_items)
+                return TypeTracerArray._new(dtype, shape=tuple(shape))
             else:
                 raise wrap_error(TypeError)
 

tests/test_2361_typetracer_asarray_nd.py

@@ -0,0 +1,65 @@
+# BSD 3-Clause License; see https://github.com/scikit-hep/awkward-1.0/blob/main/LICENSE
+
+import numpy as np
+import pytest
+
+from awkward._nplikes.typetracer import TypeTracer
+
+typetracer = TypeTracer.instance()
+
+
+def test_nd():
+    data = [
+        [
+            [1, 2, 3],
+            [4, 5, 6],
+        ],
+        [
+            [1, 2, 3],
+            [4, 5, 8],
+        ],
+    ]
+    result = typetracer.asarray(data, dtype=np.uint8)
+    assert result.dtype == np.dtype(np.uint8)
+    assert result.shape == (2, 2, 3)
+
+    # Check default size of array
+    array = np.array([1, 2, 3])
+    default_int_dtype = array.dtype
+
+    result = typetracer.asarray(data)
+    assert result.dtype == default_int_dtype
+    assert result.shape == (2, 2, 3)
+
+
+def test_nd_ragged():
+    data = [
+        [
+            [1, 2, 3],
+            [4, 5, 6],
+        ],
+        [
+            [1, 2, 3],
+            [4, 5],
+        ],
+    ]
+    with pytest.raises(ValueError, match=r"sequence at dimension .* does not match"):
+        typetracer.asarray(data)
+
+
+def test_unknown_scalar():
+    unknown_array = typetracer.asarray([0, 1, 2.0], dtype=np.dtype(np.float64))
+
+    data = [
+        [
+            [1, 2, 3],
+            [4, 5, 6],
+        ],
+        [
+            [1, 2, 3],
+            [4, 5, unknown_array[0]],
+        ],
+    ]
+    result = typetracer.asarray(data)
+    assert result.dtype == np.dtype(np.float64)
+    assert result.shape == (2, 2, 3)