Skip to content

Fix/choose trivial case#12090

Merged
jacobtomlinson merged 4 commits intodask:mainfrom
Oisin-M:fix/choose_trivial_case
Oct 6, 2025
Merged

Fix/choose trivial case#12090
jacobtomlinson merged 4 commits intodask:mainfrom
Oisin-M:fix/choose_trivial_case

Conversation

@Oisin-M
Copy link
Copy Markdown
Contributor

@Oisin-M Oisin-M commented Oct 3, 2025

This PR does a simple fix for the following bug:

>> indices = da.array([0, 0, 0, 0])
>> choices = (da.array([10., 20., 30., 40.]),)
>> print(da.choose(indices, choices).compute())
NotImplemented

The output arises from

dask/dask/array/core.py

Lines 4974 to 4977 in 1a306ba

try:
dtype = apply_infer_dtype(op, vals, {}, "elemwise", suggest_dtype=False)
except Exception:
return NotImplemented

which in turn comes from

dask/dask/array/core.py

Lines 479 to 486 in 1a306ba

args = [
(
np.ones_like(meta_from_array(x), shape=((1,) * x.ndim), dtype=x.dtype)
if is_arraylike(x)
else x
)
for x in args
]

which causes a failure since it later runs

>> np.choose(np.array([1]), np.array([1])
ValueError: invalid entry in choice array

I have removed the try-except which gave the NotImplemented output since I think it's misleading, and I've added the simplest fix I could think of, namely to use zeros_like instead of ones_like. This ends up calling np.choose(np.array([0]), np.array([0]) without any issues.

@github-actions
Copy link
Copy Markdown
Contributor

github-actions bot commented Oct 3, 2025

Unit Test Results

See test report for an extended history of previous test failures. This is useful for diagnosing flaky tests.

      9 files  ± 0        9 suites  ±0   3h 10m 14s ⏱️ +28s
 18 073 tests + 4   16 856 ✅ + 4   1 217 💤 ±0  0 ❌ ±0 
161 731 runs  +36  149 613 ✅ +37  12 118 💤  - 1  0 ❌ ±0 

Results for commit 450a42f. ± Comparison against base commit 1a306ba.

♻️ This comment has been updated with latest results.

Copy link
Copy Markdown
Member

@jacobtomlinson jacobtomlinson left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It looks like the broad exception was added in #3230 to enable mixing of dataframe and arrays for aligned element-wise operations.

You mention removing it because it's misleading, but I assume it's to handle unexpected cases when comparing the two kinds of objects. Given that it's looking for unexpected mismatches there are probably no good tests for this situation so I don't feel confident in merging this without exploring this further.

@Oisin-M
Copy link
Copy Markdown
Contributor Author

Oisin-M commented Oct 3, 2025

Hi @jacobtomlinson, thanks for the review. If it's better I can add back the try-except, but ideally it would nice to avoid catching all exceptions. That said, given it's not clear how to test I'm not sure how to proceed on that front.

@jacobtomlinson
Copy link
Copy Markdown
Member

I think the try/except is there to handle these unknown cases. If you can resolve your bug without removing it then I think that would be best.

Copy link
Copy Markdown
Member

@jacobtomlinson jacobtomlinson left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks great thanks!

@jacobtomlinson jacobtomlinson merged commit 4736a38 into dask:main Oct 6, 2025
23 of 24 checks passed
@Oisin-M Oisin-M deleted the fix/choose_trivial_case branch October 6, 2025 13:19
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

dask.array.choose() fails when the choice is trivial

2 participants