Skip to content

Fix flaky TestVmap.test_vmap_masked_scatter - #3421

Merged
zcbenz merged 1 commit into
ml-explore:mainfrom
zcbenz:fix-flaky-masked-scatter
Apr 20, 2026
Merged

Fix flaky TestVmap.test_vmap_masked_scatter#3421
zcbenz merged 1 commit into
ml-explore:mainfrom
zcbenz:fix-flaky-masked-scatter

Conversation

@zcbenz

@zcbenz zcbenz commented Apr 17, 2026

Copy link
Copy Markdown
Member

This test is flaky in CI and luckily I can reproduce it locally:

  ======================================================================
  FAIL: test_vmap_masked_scatter (test_vmap.TestVmap.test_vmap_masked_scatter)
  ----------------------------------------------------------------------
  Traceback (most recent call last):
    File "/home/runner/work/mlx/mlx/python/tests/test_vmap.py", line 900, in test_vmap_masked_scatter
      self.assertTrue(mx.array_equal(expected, out))
      ~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  AssertionError: array(False, dtype=bool) is not true
  
  ----------------------------------------------------------------------

angeloskath
angeloskath reviewed Apr 20, 2026
View reviewed changes
Comment thread mlx/backend/cuda/indexing.cpp Outdated
if (mask_flat.data<void>() != mask.data<void>()) {
encoder.add_temporary(mask_flat);
}
encoder.add_temporary(mask_flat);

@angeloskath angeloskath Apr 20, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

Hmm the previous code is actually correct in the sense that we should only add it as temporary if a copy was done and the pointer changed.

@zcbenz zcbenz Apr 20, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

Yeah I got the cause wrong, it turns out to be the implicit copy in data()..

@zcbenz
zcbenz force-pushed the fix-flaky-masked-scatter branch from 576c944 to 12c29f6 Compare April 20, 2026 06:44
angeloskath
angeloskath approved these changes Apr 20, 2026
View reviewed changes

@angeloskath angeloskath left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

Nice! That makes a lot of sense and it is a sneaky one hidden in plain sight.

@zcbenz
zcbenz merged commit b4ddf9b into ml-explore:main Apr 20, 2026
16 checks passed
@zcbenz
zcbenz deleted the fix-flaky-masked-scatter branch April 20, 2026 08:19
@BrewTestBot BrewTestBot mentioned this pull request Apr 22, 2026
Closed
1 task done
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@angeloskath angeloskath angeloskath approved these changes

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@zcbenz @angeloskath