Repository: scikit-hep/awkward

Issue:
ufuncs on UnionArray[RecordArray1, RecordArray2] hit an assertion against unknown_lengths
### Version of Awkward Array

2.1.3

### Description and code to reproduce

The following code:
```python3
import json
import awkward as ak
import numpy as np

fromjson = {
    "class": "RecordArray",
    "fields": ["muon", "jet"],
    "contents": [
        {
            "class": "ListOffsetArray",
            "offsets": "i64",
            "content": {
                "class": "RecordArray",
                "fields": ["pt", "eta", "phi", "crossref", "thing1"],
                "contents": [
                    {
                        "class": "NumpyArray",
                        "primitive": "int64",
                        "inner_shape": [],
                        "parameters": {},
                        "form_key": "muon_pt!",
                    },
                    {
                        "class": "NumpyArray",
                        "primitive": "int64",
                        "inner_shape": [],
                        "parameters": {},
                        "form_key": "muon_eta!",
                    },
                    {
                        "class": "NumpyArray",
                        "primitive": "int64",
                        "inner_shape": [],
                        "parameters": {},
                        "form_key": "muon_phi!",
                    },
                    {
                        "class": "ListOffsetArray",
                        "offsets": "i64",
                        "content": {
                            "class": "NumpyArray",
                            "primitive": "int64",
                            "inner_shape": [],
                            "parameters": {},
                            "form_key": "muon_crossref_content!",
                        },
                        "parameters": {},
                        "form_key": "muon_crossref_index!",
                    },
                    {
                        "class": "NumpyArray",
                        "primitive": "int64",
                        "inner_shape": [],
                        "parameters": {},
                        "form_key": "muon_thing1!",
                    },
                ],
                "parameters": {"__record__": "TestBehavior1"},
                "form_key": "muon_record!",
            },
            "parameters": {},
            "form_key": "muon_list_outer!",
        },
        {
            "class": "ListOffsetArray",
            "offsets": "i64",
            "content": {
                "class": "RecordArray",
                "fields": ["pt", "eta", "phi", "crossref", "thing1"],
                "contents": [
                    {
                        "class": "NumpyArray",
                        "primitive": "int64",
                        "inner_shape": [],
                        "parameters": {},
                        "form_key": "jet_pt!",
                    },
                    {
                        "class": "NumpyArray",
                        "primitive": "int64",
                        "inner_shape": [],
                        "parameters": {},
                        "form_key": "jet_eta!",
                    },
                    {
                        "class": "NumpyArray",
                        "primitive": "int64",
                        "inner_shape": [],
                        "parameters": {},
                        "form_key": "jet_phi!",
                    },
                    {
                        "class": "ListOffsetArray",
                        "offsets": "i64",
                        "content": {
                            "class": "NumpyArray",
                            "primitive": "int64",
                            "inner_shape": [],
                            "parameters": {},
                            "form_key": "jet_crossref_content!",
                        },
                        "parameters": {},
                        "form_key": "jet_crossref_index!",
                    },
                    {
                        "class": "NumpyArray",
                        "primitive": "int64",
                        "inner_shape": [],
                        "parameters": {},
                        "form_key": "jet_thing1!",
                    },
                ],
                "parameters": {"__record__": "TestBehavior2"},
                "form_key": "jet_record!",
            },
            "parameters": {},
            "form_key": "jet_list_outer!",
        },
    ],
    "parameters": {},
    "form_key": "outer!",
}


@ak.mixin_class(ak.behavior)
class TestBehavior:
    @ak.mixin_class_method(np.add, {"TestBehavior"})
    def add(self, other):
        return ak.zip(
            {
                "pt": self.pt + other.pt,
            },
            with_name="TestBehavior",
            behavior=self.behavior,
        )


@ak.mixin_class(ak.behavior)
class TestBehavior1(TestBehavior):
    pass
    
@ak.mixin_class(ak.behavior)
class TestBehavior2(TestBehavior):
    pass

form = ak.forms.from_json(json.dumps(fromjson))

ttlayout, report = ak._nplikes.typetracer.typetracer_with_report(form, forget_length=True)

ttarray = ak.Array(ttlayout, behavior=ak.behavior)

stack = ak.concatenate([ttarray.muon, ttarray.jet], axis=1)

add = stack + stack
```

fails with this trace
```
Traceback (most recent call last):
  File "/Users/lgray/miniforge3/envs/coffea-dev/lib/python3.8/site-packages/awkward/highlevel.py", line 1360, in __array_ufunc__
    return ak._connect.numpy.array_ufunc(ufunc, method, inputs, kwargs)
  File "/Users/lgray/miniforge3/envs/coffea-dev/lib/python3.8/site-packages/awkward/_connect/numpy.py", line 289, in array_ufunc
    out = ak._broadcasting.broadcast_and_apply(
  File "/Users/lgray/miniforge3/envs/coffea-dev/lib/python3.8/site-packages/awkward/_broadcasting.py", line 1065, in broadcast_and_apply
    out = apply_step(
  File "/Users/lgray/miniforge3/envs/coffea-dev/lib/python3.8/site-packages/awkward/_broadcasting.py", line 1044, in apply_step
    return continuation()
  File "/Users/lgray/miniforge3/envs/coffea-dev/lib/python3.8/site-packages/awkward/_broadcasting.py", line 1017, 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 1044, in apply_step
    return continuation()
  File "/Users/lgray/miniforge3/envs/coffea-dev/lib/python3.8/site-packages/awkward/_broadcasting.py", line 1017, in continuation
    return broadcast_any_list()
  File "/Users/lgray/miniforge3/envs/coffea-dev/lib/python3.8/site-packages/awkward/_broadcasting.py", line 609, in broadcast_any_list
    outcontent = apply_step(
  File "/Users/lgray/miniforge3/envs/coffea-dev/lib/python3.8/site-packages/awkward/_broadcasting.py", line 1044, in apply_step
    return continuation()
  File "/Users/lgray/miniforge3/envs/coffea-dev/lib/python3.8/site-packages/awkward/_broadcasting.py", line 1009, in continuation
    return broadcast_any_union()
  File "/Users/lgray/miniforge3/envs/coffea-dev/lib/python3.8/site-packages/awkward/_broadcasting.py", line 827, in broadcast_any_union
    assert length is not unknown_length
AssertionError

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "form_madness6.py", line 154, in <module>
    add = stack + stack
  File "/Users/lgray/miniforge3/envs/coffea-dev/lib/python3.8/site-packages/numpy/lib/mixins.py", line 21, in func
    return ufunc(self, other)
  File "/Users/lgray/miniforge3/envs/coffea-dev/lib/python3.8/site-packages/awkward/highlevel.py", line 1360, in __array_ufunc__
    return ak._connect.numpy.array_ufunc(ufunc, method, inputs, kwargs)
  File "/Users/lgray/miniforge3/envs/coffea-dev/lib/python3.8/site-packages/awkward/_errors.py", line 54, 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 64, in handle_exception
    raise self.decorate_exception(cls, exception)
AssertionError: 

See if this has been reported at https://github.com/scikit-hep/awkward/issues
```

@agoose77 


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
@@ -817,14 +817,14 @@ def broadcast_any_option():
     def broadcast_any_union():
         if not backend.nplike.known_data:
             # assert False
-            union_num_contents, length = [], None
+            union_num_contents = []
+            length = None
             for x in contents:
                 if x.is_union:
                     x._touch_data(recursive=False)
                     union_num_contents.append(len(x.contents))
                     if length is None:
-                        length = x.tags.data.shape[0]
-            assert length is not unknown_length
+                        length = x.length
 
             all_combos = list(
                 itertools.product(*[range(x) for x in union_num_contents])
@@ -832,7 +832,8 @@ def broadcast_any_union():
 
             tags = backend.index_nplike.empty(length, dtype=np.int8)
             index = backend.index_nplike.empty(length, dtype=np.int64)
-            numoutputs, outcontents = None, []
+            numoutputs = None
+            outcontents = []
             for combo in all_combos:
                 nextinputs = []
                 i = 0
@@ -856,13 +857,9 @@ def broadcast_any_union():
                     )
                 )
                 assert isinstance(outcontents[-1], tuple)
-                if numoutputs is None:
-                    numoutputs = outcontents[-1].length
-                else:
-                    assert (
-                        numoutputs is unknown_length
-                        or outcontents[-1].length is unknown_length
-                    ) or numoutputs == outcontents[-1].length
+                if numoutputs is not None:
+                    assert numoutputs == len(outcontents[-1])
+                numoutputs = len(outcontents[-1])
 
             assert numoutputs is not None
 


Test patch:
diff --git a/tests/test_2418_union_broadcast_unknown.py b/tests/test_2418_union_broadcast_unknown.py
new file mode 100644
--- /dev/null
+++ b/tests/test_2418_union_broadcast_unknown.py
@@ -0,0 +1,14 @@
+# BSD 3-Clause License; see https://github.com/scikit-hep/awkward-1.0/blob/main/LICENSE
+
+
+import awkward as ak
+
+
+def test():
+    array = ak.Array([1, [2, 3]])
+    result = array + array
+
+    array_tt = ak.Array(array.layout.to_typetracer(forget_length=True))
+    result_tt = array_tt + array_tt
+
+    assert result_tt.layout.form == result.layout.form