awkward's named axes causes regression in test_notebooks

[veprbl]

Vector Version

1.5.2

Python Version

3.11 and 3.12

OS / Environment

nixpkgs
NixOS/nixpkgs#354346

Describe the bug

Failure in tests/test_notebooks.py

E           papermill.exceptions.PapermillExecutionError: 
E           ---------------------------------------------------------------------------
E           Exception encountered at "In [96]":
E           ---------------------------------------------------------------------------
E           AttributeError                            Traceback (most recent call last)
E           Cell In[96], line 2
E                 1 (
E           ----> 2     vector.Array(
E                 3         [  # an Awkward Array of vectors
E                 4             [{"x": 1, "y": 1.1}, {"x": 2, "y": 2.2}],
E                 5             [],
E                 6             [{"x": 3, "y": 3.3}],
E                 7             [{"x": 4, "y": 4.4}, {"x": 5, "y": 5.5}],
E                 8         ]
E                 9     )
E                10     + vector.obj(x=10, y=5)  # and a single vector object
E                11 )
E           
E           File /nix/store/yq4a8j78af4l3rcc1m7sf5kqqv3ydnkb-python3.12-awkward-2.7.0/lib/python3.12/site-packages/awkward/_operators.py:54, in _binary_method.<locals>.func(self, other)
E                51 if _disables_array_ufunc(other):
E                52     return NotImplemented
E           ---> 54 return ufunc(self, other)
E           
E           File /nix/store/yq4a8j78af4l3rcc1m7sf5kqqv3ydnkb-python3.12-awkward-2.7.0/lib/python3.12/site-packages/awkward/highlevel.py:1594, in Array.__array_ufunc__(self, ufunc, method, *inputs, **kwargs)
E              1592 name = f"{type(ufunc).__module__}.{ufunc.__name__}.{method!s}"
E              1593 with ak._errors.OperationErrorContext(name, inputs, kwargs):
E           -> 1594     return ak._connect.numpy.array_ufunc(ufunc, method, inputs, kwargs)
E           
E           File /nix/store/yq4a8j78af4l3rcc1m7sf5kqqv3ydnkb-python3.12-awkward-2.7.0/lib/python3.12/site-packages/awkward/_connect/numpy.py:367, in array_ufunc(ufunc, method, inputs, kwargs)
E               364 attrs = attrs_of(*inputs)
E               365 backend = backend_of(*inputs, coerce_to_common=True)
E           --> 367 depth_context, lateral_context = NamedAxesWithDims.prepare_contexts(inputs)
E               369 inputs = _array_ufunc_custom_cast(inputs, behavior, backend)
E               371 def action(inputs, **ignore):
E           
E           File /nix/store/yq4a8j78af4l3rcc1m7sf5kqqv3ydnkb-python3.12-awkward-2.7.0/lib/python3.12/site-packages/awkward/_namedaxis.py:686, in NamedAxesWithDims.prepare_contexts(cls, arrays, unwrap_kwargs)
E               684         layout = ctx.unwrap(array, **_unwrap_kwargs)
E               685     _named_axes.append(_get_named_axis(array))
E           --> 686     _ndims.append(layout.minmax_depth[1])
E               688 depth_context = {NAMED_AXIS_KEY: cls(_named_axes, _ndims)}
E               689 lateral_context = {NAMED_AXIS_KEY: cls(_named_axes, _ndims)}
E           
E           AttributeError: 'VectorObject2D' object has no attribute 'minmax_depth'
E           
E           This error occurred while calling
E           
E               numpy.add.__call__(
E                   <VectorArray2D [[{x: 1, y: 1.1}, {...}], ...] type='4 * var * Vecto...'>
E                   VectorObject2D-instance
E               )

Any additional but relevant log output

No response

[henryiii]

CC @pfackeldey

[pfackeldey]

Hm, this is because ak.to_layout(vector.obj(x=10, y=5), allow_unknown) will return vector.obj(x=10, y=5) again (and not an awkward layout type). And this object does not describe its number of dimensions as a layout does. But this is needed here for broadcasting with named axis.

Currently, I see 2 options to fix this:

1. ak.to_layout(vector.obj(x=10, y=5), allow_unknown) would transform into an awkward layout that describes its depth correctly, similar to how ak.to_layout(5) returns <NumpyArray dtype='int64' len='1'>[5]</NumpyArray>.
2. awkward-array checks if it can correctly to_layout the inputs for broadcasting (and if none of the inputs has named axis attached), and if (both) not: circumvent the whole named_axis logic. I'm not yet sure how to do this and if this is sufficient, I'll have a look.

What do you think @jpivarski ?

I'm not sure what the correct policy would be for awkward-array.
My understanding was/is that awkward-array works with anything correctly that ak.to_layout understands. If it doesn't there's no guarantee that things work with awkward-array correctly. Is this wrong?

PS: I think there might be a "simple" solution by making a vector.obj iterable, i.e. adding __iter__ on it. Then, it could be interpreted by the ak.to_layout implementation as an awkward array with a layout... Nevermind, just tested that: in principle this works, but it's lacking the correct __record__ parameter in the layout then.

[jpivarski]

The vector.obj should not be iterable.

But I don't see why adding named axis capabilities should cause anything that used to work to not work. The example in the notebook has an array of vectors (ak.Array) plus a single vector (vector.obj), and these are passed into the broadcasting and ufunc-handling logic with the single vector not expanded into an array, primarily for performance but on last Friday's call, Martin found it useful that the fact that it was scalar (what he called "literal") is preserved.

Since this argument does not have ak.Array (or ak.Record) type, the named axis logic should not try to determine what the names of its axes are. Instead of having two sets of axis names to merge, this case has only one set of axis names. The function argument could have been an array and could have had axis names, but did not. The axis name propagation should allow for that. I think this is your proposal number 2.

Do any of your name-merging strategies allow for an empty set of axis names to act like a zero in addition? That is, if you have a non-unary function and think you might have to merge N sets of axis names, an empty set means "nothing to merge; take the other set of names as-is"? If so, that might be an easy way to implement the not isinstance(ak.to_layout(·), ak.Array) case.

[pfackeldey]

Since this argument does not have ak.Array (or ak.Record) type, the named axis logic should not try to determine what the names of its axes are. Instead of having two sets of axis names to merge, this case has only one set of axis names. The function argument could have been an array and could have had axis names, but did not. The axis name propagation should allow for that. I think this is your proposal number 2.

Yes, this is my proposal number 2.
This is not about the incorrectly inferring named axis (they'll be inferred correctly as empty from any object if they have non attached), but rather about the number of "initial" dimensions (before broadcasting) that I need to propagate named axis correctly through broadcasting. I assumed that only objects that can be interpreted with ak.to_layout are broadcastable together (this assumption allowed me to get the number of "initial" dimensions from the layouts). It sounds like my assumption here was wrong, which is why I need to soften this in awkward's broadcasting logic <- proposal number 2.

Do any of your name-merging strategies allow for an empty set of axis names to act like a zero in addition? That is, if you have a non-unary function and think you might have to merge N sets of axis names, an empty set means "nothing to merge; take the other set of names as-is"? If so, that might be an easy way to implement the not isinstance(ak.to_layout(·), ak.Array) case.

Yes, this is what happens with 2 arrays where only one has named axis.

[pfackeldey]

Hi @veprbl,
can you test if PR scikit-hep/awkward#3299 fixes this issue for you?