Proposed changes

Follow for #2860, see ASAN issue in a PR run where we had the new CI check enabled: https://github.com/ml-explore/mlx/actions/runs/19943012002/job/57185456349?pr=2860

My understanding of the logic I am attempting to improve:
(@awni please call me out if I got it wrong)

std::tie(ptr[count.first], ptr[count.second]) =
            random::threefry2x32_hash(key, count);

… is filling the buffer from both ends inward (threefry2x32_hash always returns a pair).

The loop stops just before the pointers from the start and end would cross over or meet. Therefore, the two if statements after it are for handling the “leftovers”.

The first if statement (if (count.first < half_size) {) handles the case where we have 2 slots empty in the middle, e.g. [ X, X, X, _, _, X, X, X ].

both rb.first and rb.second are used because we are filling the middle pair of slots.

The second if statement (if (!even) {) is where ASan triggered.

There is an odd number of slots and only one is left to fill in the middle, e.g. [ X, _, X].
Here we need to generate a new final random number on the fly and use it to fill that unique middle slot (safely).

Proposed Remediation:

The proposed fix aims to remove unsafe, unconditional 4-byte writes. Got rid of direct assignment -- now first confirming if the write position is even valid, then calculating bytes remaining in the buffer and copying only what fits (ideally full 4 bytes, but less if it wouldn’t fit).

Locally I tested it with the ASan and it seems to be fixed now.

CC @madrob

Checklist

Put an x in the boxes that apply.

• I have read the CONTRIBUTING document
• I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
• I have added tests that prove my fix is effective or that my feature works
• I have updated the necessary documentation (if needed)