[BUG] VJP should be Sendable

[davidkoski]

Two parter, see ml-explore/mlx-swift-lm#616

public func CustomFunction(
    @MLXCustomFunctionBuilder.Builder _ build: () -> ([MLXArray]) -> [MLXArray]
) -> ([MLXArray]) -> [MLXArray] {
    build()
}

should produce Sendable ouputs -- these are pure functions.

Also VJP:

public func VJP(_ f: @escaping ([MLXArray], [MLXArray]) -> [MLXArray]) -> MLXCustomFunctionComponent
{
    .vjp(f)
}

should have a 3 parameter mode with outputs, see:

    std::optional<std::function<std::vector<array>(
        const std::vector<array>&,
        const std::vector<array>&,
        const std::vector<array>&)>> fun_vjp /* = std::nullopt */,

this can be used with:

public func depends(input: MLXArray, dependencies: [MLXArray]) -> MLXArray {

[shubhransh-gupta]

Hi @davidkoski,

I'd like to pick this up!

Here is the proposed technical approach:

1. Sendable Support

• Mark CustomFunction and buildBlock return types as @Sendable ([MLXArray]) -> [MLXArray]:

  public func CustomFunction(
      @MLXCustomFunctionBuilder.Builder _ build: () -> @Sendable ([MLXArray]) -> [MLXArray]
  ) -> @Sendable ([MLXArray]) -> [MLXArray]

• Mark component builder functions and closure parameters as @Sendable:
  • Forward(_ f: @escaping @Sendable ([MLXArray]) -> [MLXArray]) -> MLXCustomFunctionComponent
  • VJP(_ f: @escaping @Sendable ([MLXArray], [MLXArray]) -> [MLXArray]) -> MLXCustomFunctionComponent
• _CustomFunctionState is already @unchecked Sendable (thread-safe via internal NSLock), so the returned closure can safely be @Sendable.

2. 3-Parameter VJP Support (primals, cotangents, outputs)

• In _CustomFunctionState.buildClosures(), mlx_closure_custom_new_func_payload already receives input_2 (outputs) in C (Source/Cmlx/mlx-c/mlx/c/closure.cpp), but the 4th callback parameter was currently discarded (_). We will unpack mlx_vector_array_values(outputs) and forward it.
• Add an overload for VJP accepting (primals, cotangents, outputs):

  public func VJP(
      _ f: @escaping @Sendable (_ primals: [MLXArray], _ cotangents: [MLXArray], _ outputs: [MLXArray]) -> [MLXArray]
  ) -> MLXCustomFunctionComponent

• For MLXCustomFunctionComponent:
  • We can either keep .vjp(@Sendable ([MLXArray], [MLXArray], [MLXArray]) -> [MLXArray]) internally (where the 2-arg VJP adapts as { p, c, _ in f(p, c) }), or maintain explicit enum cases (e.g. .vjp and .vjp3 / .vjpWithOutputs). Do you have a preference here regarding public enum casing?

3. Testing & Verification

• Add unit tests verifying:
  • 3-parameter VJP receiving forward outputs (e.g., verifying recurrence / checkpointing pattern with depends(inputs:dependencies:)).
  • Sendability of CustomFunction outputs across Task / concurrent boundaries without compiler warnings.
  • Backward compatibility for existing 2-parameter VJP callers.

Please let me know if this approach looks good to you or if you'd like any adjustments before I submit the PR!

[davidkoski]

This looks right to me. See also #469 -- I ended up making the interior be Sendable and the callers would return Sendable or not depending on the incoming closure. I am not sure VJP needs a non-Sendable version. I think it should be a pure function.

[shubhransh-gupta]

Hi @davidkoski, thanks for the feedback! I've opened PR #481 with these changes:

• Forward and VJP closures are now @Sendable, and CustomFunction returns @Sendable ([MLXArray]) -> [MLXArray].
• Added the 3-parameter overload VJP(_ f: @escaping @Sendable (_ primals: [MLXArray], _ cotangents: [MLXArray], _ outputs: [MLXArray]) -> [MLXArray]).
• Added test coverage in TransformTests for both 2-param and 3-param VJP, as well as cross-task Sendability.