Skip to content

Add Sendable closures and 3-parameter VJP support to MLXCustomFunction - #481

Open
shubhransh-gupta wants to merge 1 commit into
ml-explore:mainfrom
shubhransh-gupta:vjp-sendable-three-params
Open

shubhransh-gupta wants to merge 1 commit into
ml-explore:mainfrom
shubhransh-gupta:vjp-sendable-three-params

Conversation

@shubhransh-gupta

Copy link
Copy Markdown

Description

Following discussion on #475:

  1. Make Forward and VJP components conform to @Sendable, ensuring custom functions are pure and safe across concurrency contexts.
  2. Make CustomFunction return @Sendable ([MLXArray]) -> [MLXArray].
  3. Add a 3-parameter overload for VJP:
    public func VJP(
        _ f: @escaping @Sendable (_ primals: [MLXArray], _ cotangents: [MLXArray], _ outputs: [MLXArray]) -> [MLXArray]
    ) -> MLXCustomFunctionComponent
    This passes through the forward outputs already provided by the underlying C closure (mlx_closure_custom_new_func_payload) without needing extra C++/Cmlx changes.
  4. Maintain full backward compatibility for the 2-parameter VJP (_ primals: [MLXArray], _ cotangents: [MLXArray]) signature.
  5. Add unit tests in TransformTests covering 2-parameter VJP, 3-parameter VJP using outputs, and cross-task Sendability.

Closes #475

@davidkoski

Copy link
Copy Markdown
Member

This needs a swift-format run (the lint failure)

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.

[BUG] VJP should be Sendable

2 participants