Skip to content

Keep CompiledTrace's compiled function non-Sendable (fixes Sendable mlx-swift compile) #87

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
solderzzc merged 1 commit into from
Oct 21, 2026
Merged

Keep CompiledTrace's compiled function non-Sendable (fixes Sendable mlx-swift compile) #87

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 9 additions & 18 deletions Libraries/MLXLMCommon/CompiledTrace.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -99,7 +99,7 @@ public final class CompiledTrace<Owner: Module>: CompiledTraceInvalidating {
// Compiled on first call rather than at init, because weights load later.
// The lock keeps two concurrent first calls from installing rival traces.
private let lock = NSLock()
private var function: (@Sendable ([MLXArray]) -> [MLXArray])?
private var function: (([MLXArray]) -> [MLXArray])?
private var tracedOwner: ObjectIdentifier?

/// - Parameters:
Expand Down Expand Up @@ -146,7 +146,7 @@ public final class CompiledTrace<Owner: Module>: CompiledTraceInvalidating {
self(owner, [argument])[0]
}

private func compiled(_ owner: Owner) -> @Sendable ([MLXArray]) -> [MLXArray] {
private func compiled(_ owner: Owner) -> ([MLXArray]) -> [MLXArray] {
let body = self.body

// A snapshot is enough: `update(parameters:)`, the path optimizers and
Expand All @@ -158,26 +158,17 @@ public final class CompiledTrace<Owner: Module>: CompiledTraceInvalidating {

// The one place in this package that captures a module in a compiled
// body: `unowned` because the module stores the trace, and a strong
// capture would cycle and leak the weights with it. The capture goes
// through a box because `MLX.compile` takes a `@Sendable` closure and a
// module is not `Sendable`.
let owned = UnownedOwner(owner)
return MLX.compile(inputs: [weights]) { arguments in
body(owned.owner, arguments)
// capture would cycle and leak the weights with it.
//
// `MLX.compile(inputs:)` returns a plain (non-`@Sendable`) function, so
// the trace keeps that type too. Declaring it `@Sendable` fails to compile
// against mlx-swift's `compile` overloads that take `@Sendable` closures.
return MLX.compile(inputs: [weights]) { [unowned owner] arguments in
body(owner, arguments)
}
}
}

/// Holds a module unowned so a compiled closure can name it without retaining it.
/// Unchecked: the trace is stored by that module and only ever runs against it.
private final class UnownedOwner<Owner: Module>: @unchecked Sendable {
unowned let owner: Owner

init(_ owner: Owner) {
self.owner = owner
}
}

/// Something `Module.invalidateCompiledTraces()` can reset: a ``CompiledTrace``
/// or a container of them.
public protocol CompiledTraceInvalidating: AnyObject {
Expand Down
Loading