Convert MLX shared variable updates back to numpy arrays - #2378
Open
jessegrabowski wants to merge 4 commits into
Open
Convert MLX shared variable updates back to numpy arrays#2378jessegrabowski wants to merge 4 commits into
jessegrabowski wants to merge 4 commits into
Conversation
Copy link
Copy Markdown
Member
Sorry, something went wrong.
The value written into a shared variable's container outlives the call and may later be read by a function compiled for another backend, so JITLinker.output_filter is applied to update outputs while returned outputs keep their native backend type.
jessegrabowski
requested review from
cetagostini,
juanitorduz,
ricardoV94 and
williambdean
August 25, 2026 00:24
Copy link
Copy Markdown
Member
Author
|
@juanitorduz tagged you because i was thinking about #2375, wondering if the solution to that wasn't also similar to this, adjustment of the linker behavior in a certain case for a certain backend. |
Sorry, something went wrong.
Copy link
Copy Markdown
Member
|
Check out #2285 You don't always care about cross backend compatibility so it offers an escape. It's also lazy, only done when a different backend would actually read it. Still worrried about overhead though. |
Sorry, something went wrong.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Reviewers
At least 1 approving review is required to merge this pull request.
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Function.__call__ writes an update straight into the shared variable's storage slot, bypassing the container's filter, so after one MLX call the shared variable holds an mx.array instead of an ndarray. get_value() then returns a device array even with borrow=False, and passing that shared variable to a function compiled for another backend fails outright - a plain numba graph raises a TypingError on it.
JITLinker.output_filter already existed for this and was dead code, called by nothing. It's now applied to the outputs named in fgraph.update_mapping, and MLXLinker overrides it. Values returned to the caller still come back as mx.array; only what gets stored back in a container is converted. Backends that don't override the hook run the same thunk as before.