Repository: scikit-hep/awkward

Issue:
Assigning `[]` or `{}` into RecordArray typetracer fails
### Version of Awkward Array

2.1.4

### Description and code to reproduce

The following script:
```python3
import json
import awkward as ak
import numpy as np

fromjson = {
    "class": "RecordArray",
    "fields": ["muon"],
    "contents": [
        {
            "class": "ListOffsetArray",
            "offsets": "i64",
            "content": {
                "class": "RecordArray",
                "fields": ["pt"],
                "contents": [
                    {
                        "class": "NumpyArray",
                        "primitive": "int64",
                        "inner_shape": [],
                        "parameters": {},
                        "form_key": "muon_pt!",
                    },
                ],
                "parameters": {},
                "form_key": "muon_record!",
            },
            "parameters": {},
            "form_key": "muon_list_outer!",
        },
    ],
    "parameters": {},
    "form_key": "outer!",
}

form = ak.forms.from_json(json.dumps(fromjson))

ttlayout, report = ak._nplikes.typetracer.typetracer_with_report(form, forget_length=True)

ttarray = ak.Array(ttlayout)

ttmuons = ttarray.muon

array = ak.Array(ttlayout.form.length_zero_array(highlevel=False))

array["emptydict"] = {}

muons = array.muon

muons["emptydict"] = {}

array["emptylist"] = []

muons["emptylist"] = []

# all four of the following fail

ttarray["emptydict"] = {}

ttmuons["emptydict"] = {}

ttarray["emptylist"] = []

ttmuons["emptylist"] = []
```
results in (and similar variants for different types being assigned):
```
Traceback (most recent call last):
  File "/Users/lgray/miniforge3/envs/coffea-dev/lib/python3.8/site-packages/awkward/highlevel.py", line 1027, in __setitem__
    self._layout = ak.operations.with_field(
  File "/Users/lgray/miniforge3/envs/coffea-dev/lib/python3.8/site-packages/awkward/operations/ak_with_field.py", line 47, in with_field
    return _impl(array, what, where, highlevel, behavior)
  File "/Users/lgray/miniforge3/envs/coffea-dev/lib/python3.8/site-packages/awkward/operations/ak_with_field.py", line 153, in _impl
    out = ak._broadcasting.broadcast_and_apply(
  File "/Users/lgray/miniforge3/envs/coffea-dev/lib/python3.8/site-packages/awkward/_broadcasting.py", line 1062, in broadcast_and_apply
    out = apply_step(
  File "/Users/lgray/miniforge3/envs/coffea-dev/lib/python3.8/site-packages/awkward/_broadcasting.py", line 1041, in apply_step
    return continuation()
  File "/Users/lgray/miniforge3/envs/coffea-dev/lib/python3.8/site-packages/awkward/_broadcasting.py", line 1014, in continuation
    return broadcast_any_list()
  File "/Users/lgray/miniforge3/envs/coffea-dev/lib/python3.8/site-packages/awkward/_broadcasting.py", line 563, in broadcast_any_list
    outcontent = apply_step(
  File "/Users/lgray/miniforge3/envs/coffea-dev/lib/python3.8/site-packages/awkward/_broadcasting.py", line 1027, in apply_step
    result = action(
  File "/Users/lgray/miniforge3/envs/coffea-dev/lib/python3.8/site-packages/awkward/operations/ak_with_field.py", line 144, in action
    out = ak.contents.RecordArray(
  File "/Users/lgray/miniforge3/envs/coffea-dev/lib/python3.8/site-packages/awkward/contents/recordarray.py", line 161, in __init__
    raise TypeError(
TypeError: RecordArray 'contents' must use the same array library (backend): TypeTracerBackend vs NumpyBackend

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "form_madness12.py", line 59, in <module>
    ttarray["emptylist"] = []
  File "/Users/lgray/miniforge3/envs/coffea-dev/lib/python3.8/site-packages/awkward/highlevel.py", line 1030, in __setitem__
    self._numbaview = None
  File "/Users/lgray/miniforge3/envs/coffea-dev/lib/python3.8/site-packages/awkward/_errors.py", line 56, in __exit__
    self.handle_exception(exception_type, exception_value)
  File "/Users/lgray/miniforge3/envs/coffea-dev/lib/python3.8/site-packages/awkward/_errors.py", line 66, in handle_exception
    raise self.decorate_exception(cls, exception)
TypeError: RecordArray 'contents' must use the same array library (backend): TypeTracerBackend vs NumpyBackend

This error occurred while calling

    ak.Array.__setitem__(
        self = <Array-typetracer [...] type='## * {muon: var * {pt: int64}}'>
        field_name = 'emptylist'
        field_value = []
    )
```


Discussion:

Resolving patch:
diff --git a/src/awkward/_broadcasting.py b/src/awkward/_broadcasting.py
--- a/src/awkward/_broadcasting.py
+++ b/src/awkward/_broadcasting.py
@@ -67,6 +67,7 @@ def length_of_broadcast(inputs: Sequence) -> int | type[unknown_length]:
 def broadcast_pack(inputs: Sequence, isscalar: list[bool]) -> list:
     maxlen = length_of_broadcast(inputs)
     nextinputs = []
+
     for x in inputs:
         if isinstance(x, Record):
             index = x.backend.index_nplike.full(maxlen, x.at, dtype=np.int64)
@@ -76,7 +77,7 @@ def broadcast_pack(inputs: Sequence, isscalar: list[bool]) -> list:
             nextinputs.append(
                 RegularArray(
                     x,
-                    x.length if x.backend.nplike.known_data else 1,
+                    x.length,
                     1,
                     parameters=None,
                 )
@@ -497,68 +498,100 @@ def broadcast_any_record():
         )
 
     def broadcast_any_list():
+        index_nplike = backend.index_nplike
         # All regular?
         if all(x.is_regular or not x.is_list for x in contents):
             # Ensure all layouts have same length
-            length = unset
+            length = None
             for x in contents:
-                if length is unset:
+                if length is None:
                     length = x.length
-                elif backend.nplike.known_data:
+                elif length is not unknown_length and x.length is not unknown_length:
                     assert length == x.length
-            assert length is not unset
+            assert length is not None
 
-            if any(x.size == 0 for x in contents if x.is_regular):
-                dimsize = 0
-            else:
-                dimsize = max(x.size for x in contents if x.is_regular)
-
-            if backend.nplike.known_data:
-                for x in contents:
-                    if x.is_regular:
-                        if dimsize > 1 and x.size == 1:
-                            # For any (N, 1) array, we know we'll broadcast to (N, M) where M is maxsize
-                            tmpindex = Index64(
-                                backend.index_nplike.repeat(
-                                    backend.index_nplike.arange(
-                                        x.length, dtype=np.int64
-                                    ),
-                                    dimsize,
+            # Determine the size of the broadcast result
+            dim_size = None
+            for x in contents:
+                if not x.is_regular:
+                    continue
+
+                # Any unknown length sets max_size to unknown
+                if x.size is unknown_length:
+                    dim_size = unknown_length
+                # Any zero-length column triggers zero broadcasting
+                elif x.size == 0:
+                    dim_size = 0
+                    break
+                elif dim_size is None:
+                    dim_size = x.size
+                elif dim_size is unknown_length:
+                    continue
+                else:
+                    dim_size = max(dim_size, x.size)
+
+            dimsize_greater_than_one_if_known = (
+                dim_size is unknown_length or dim_size > 1
+            )
+            dimsize_known_to_be_zero = dim_size is not unknown_length and dim_size == 0
+
+            # Build a broadcast index for size=1 contents
+            size_one_carry_index = None
+            for x in contents:
+                if x.is_regular:
+                    x_size_known_to_be_one = (
+                        x.size is not unknown_length and x.size == 1
+                    )
+                    if dimsize_greater_than_one_if_known and x_size_known_to_be_one:
+                        # For any (N, 1) array, we know we'll broadcast to (N, M) where M is maxsize
+                        size_one_carry_index = Index64(
+                            index_nplike.repeat(
+                                index_nplike.arange(
+                                    index_nplike.shape_item_as_index(x.length),
+                                    dtype=np.int64,
                                 ),
-                                nplike=backend.index_nplike,
-                            )
+                                index_nplike.shape_item_as_index(dim_size),
+                            ),
+                            nplike=index_nplike,
+                        )
+                        break
 
-                nextinputs = []
-                for x in inputs:
-                    if isinstance(x, RegularArray):
-                        if dimsize > 1 and x.size == 1:
-                            nextinputs.append(
-                                x.content[: x.length * x.size]._carry(
-                                    tmpindex, allow_lazy=False
-                                )
-                            )
-                        elif x.size == dimsize:
-                            nextinputs.append(x.content[: x.length * x.size])
-                        elif dimsize == 0:
-                            nextinputs.append(x.content[:0])
-                        else:
-                            raise ValueError(
-                                "cannot broadcast RegularArray of size "
-                                "{} with RegularArray of size {} {}".format(
-                                    x.size, dimsize, in_function(options)
-                                )
+            # Here we have three possible broadcasting outcomes (by precedence):
+            # 1. any (exactly) size-0 content trims all other contents
+            # 2. any (exactly) size-1 content broadcasts to the common length
+            # 3. otherwise, recurse into the content as-is
+            nextinputs = []
+            for x in inputs:
+                if isinstance(x, RegularArray):
+                    x_size_known_to_be_one = (
+                        x.size is not unknown_length and x.size == 1
+                    )
+                    # If dimsize is known to be exactly zero, all contents are zero length
+                    if dimsize_known_to_be_zero:
+                        nextinputs.append(x.content[:0])
+                    # If we have a known size=1 content, then broadcast it to the dimension size
+                    elif dimsize_greater_than_one_if_known and x_size_known_to_be_one:
+                        nextinputs.append(
+                            x.content[: x.length * x.size]._carry(
+                                size_one_carry_index, allow_lazy=False
                             )
+                        )
+                    # Any unknown values or sizes are assumed to be correct as-is
+                    elif (
+                        dim_size is unknown_length
+                        or x.size is unknown_length
+                        or x.size == dim_size
+                    ):
+                        nextinputs.append(x.content[: x.length * x.size])
                     else:
-                        nextinputs.append(x)
-
-            else:
-                nextinputs = []
-                for x in inputs:
-                    if isinstance(x, RegularArray):
-                        x._touch_data(recursive=False)
-                        nextinputs.append(x.content)
-                    else:
-                        nextinputs.append(x)
+                        raise ValueError(
+                            "cannot broadcast RegularArray of size "
+                            "{} with RegularArray of size {} {}".format(
+                                x.size, dim_size, in_function(options)
+                            )
+                        )
+                else:
+                    nextinputs.append(x)
 
             outcontent = apply_step(
                 backend,
@@ -573,7 +606,7 @@ def broadcast_any_list():
             assert isinstance(outcontent, tuple)
             parameters = parameters_factory(len(outcontent))
             return tuple(
-                RegularArray(x, dimsize, length, parameters=p)
+                RegularArray(x, dim_size, length, parameters=p)
                 for x, p in zip(outcontent, parameters)
             )
 
@@ -600,7 +633,6 @@ def broadcast_any_list():
                     )
                     nextinputs.append(x.content)
                 elif isinstance(x, RegularArray):
-                    x._touch_data(recursive=False)
                     nextinputs.append(x.content)
                 else:
                     nextinputs.append(x)
diff --git a/src/awkward/operations/ak_with_field.py b/src/awkward/operations/ak_with_field.py
--- a/src/awkward/operations/ak_with_field.py
+++ b/src/awkward/operations/ak_with_field.py
@@ -3,6 +3,8 @@
 import copy
 
 import awkward as ak
+from awkward._backends.dispatch import backend_of
+from awkward._backends.numpy import NumpyBackend
 from awkward._behavior import behavior_of
 from awkward._layout import wrap_layout
 from awkward._nplikes.numpylike import NumpyMetadata
@@ -11,6 +13,9 @@
 np = NumpyMetadata.instance()
 
 
+cpu = NumpyBackend.instance()
+
+
 def with_field(array, what, where=None, *, highlevel=True, behavior=None):
     """
     Args:
@@ -81,8 +86,14 @@ def _impl(base, what, where, highlevel, behavior):
             where = where[0]
 
         behavior = behavior_of(base, what, behavior=behavior)
-        base = ak.operations.to_layout(base, allow_record=True, allow_other=False)
+        backend = backend_of(base, what, default=cpu)
+
+        base = ak.operations.to_layout(
+            base, allow_record=True, allow_other=False
+        ).to_backend(backend)
         what = ak.operations.to_layout(what, allow_record=True, allow_other=True)
+        if isinstance(what, (ak.contents.Content, ak.record.Record)):
+            what = what.to_backend(backend)
 
         keys = copy.copy(base.fields)
         if where in base.fields:


Test patch:
diff --git a/tests/test_2444_minimal_listarray.py b/tests/test_2444_minimal_listarray.py
new file mode 100644
--- /dev/null
+++ b/tests/test_2444_minimal_listarray.py
@@ -0,0 +1,21 @@
+import awkward as ak
+
+
+def test_dict():
+    array = ak.Array([{"muon": [{"pt": 1}, {"pt": 4}]}])
+    ttarray = ak.to_backend(array, "typetracer")
+
+    array["emptydict"] = {}
+    ttarray["emptydict"] = {}
+
+    assert ttarray.layout.form == array.layout.form
+
+
+def test_list():
+    array = ak.Array([{"muon": [{"pt": 1}, {"pt": 4}]}])
+    ttarray = ak.to_backend(array, "typetracer")
+
+    array["emptylist"] = []
+    ttarray["emptylist"] = []
+
+    assert ttarray.layout.form == array.layout.form