fix: multiple-output ufuncs e.g. divmod
#2654
Conversation
agoose77
changed the title
divmod
|
@jpivarski in this PR so-far, I removed the unary pathway for ufuncs. I can't see why this was introduced in 0e0202a It might just be that our broadcasting code has changed since then to make the need redundant. Are you OK with the removal? |
Codecov Report
Additional details and impacted files
|
There was a problem hiding this comment.
Thanks for catching this! At one point, the return value for broadcasting actions had to be promoted from a single Content to a tuple of Contents, and I think it was because of ufuncs that return multiple values. But when that went it, it should have gone in with a test, like this one.
If you're done with this PR, go ahead and merge it.
| if sum(int(isinstance(x, ak.contents.Content)) for x in inputs) == 1: | ||
| where = None | ||
| for i, x in enumerate(inputs): | ||
| if isinstance(x, ak.contents.Content): | ||
| where = i | ||
| break | ||
| assert where is not None | ||
|
|
||
| nextinputs = list(inputs) | ||
|
|
||
| def unary_action(layout, **ignore): | ||
| nextinputs[where] = layout | ||
| result = action(tuple(nextinputs), **ignore) | ||
| if result is None: | ||
| return None | ||
| else: | ||
| assert isinstance(result, tuple) and len(result) == 1 | ||
| return result[0] | ||
|
|
||
| out = ak._do.recursively_apply( | ||
| inputs[where], | ||
| unary_action, | ||
| behavior, | ||
| function_name=ufunc.__name__, | ||
| allow_records=False, | ||
| ) |
There was a problem hiding this comment.
@jpivarski in this PR so-far, I removed the unary pathway for ufuncs. I can't see why this was introduced in 0e0202a
That commit is a PR merge; for finer granularity, the commit within that PR is
5b0a273#diff-e7b6632c2ca4228bf21505ced52be60d4fda8813e00cbaf65cd1c5d2e228d721R242-R263
It was accompanied by a large batch of single-line assertions, some of which have only one Content in the ufunc (such as + with one Content and one Python number).
Broadcasting a single layout should be equivalent to recursively_apply on that one layout. The fact that numerous tests were probably going through that and they passed either way is good evidence that the "should be" holds. (I don't see anything special about it, that it has to be there, and the context of the commit isn't reminding me of any gotchas.)
It turns out that we didn't yet test the
nout != 1ufuncs, such asdivmod.This PR adds support for this ufunc, with some refactoring work along the way.