Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
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
4 changes: 1 addition & 3 deletions Source/MLX/MLXArray.swift
Original file line number Diff line number Diff line change
Expand Up @@ -554,9 +554,7 @@ public final class MLXArray {
/// MLX is lazy and arrays are not fully realized until they are evaluated. This method is typically
/// not needed as all reads ensure the contents are evaluated.
public func eval() {
_ = withEvalLock {
mlx_array_eval(ctx)
}
MLX.eval(self)
}

/// Replace the contents with a reference to a new array (INTERNAL).
Expand Down
25 changes: 23 additions & 2 deletions Source/MLX/Transforms+Eval.swift
Original file line number Diff line number Diff line change
Expand Up @@ -115,13 +115,31 @@ func withEvalLock<R>(_ body: () throws -> R) rethrows -> R {
return try body()
}

/// Is the array's data computed?
///
/// This is the observable difference between ``eval(_:)-(MLXArray...)`` and
/// ``asyncEval(_:)-(Collection<MLXArray>)``, and the only direct way to check
/// that `eval` is still synchronous now that it schedules with
/// `mlx_async_eval` and waits with `mlx_eval`.
///
/// Internal: exists for the tests (`EvalTests`), which cannot reach `Cmlx`
/// directly. Call this only on the thread that owns `array`.
func isEvaluated(_ array: MLXArray) -> Bool {
var available = false
_mlx_array_is_available(&available, array.ctx)
return available
}

/// Evaluate one or more `MLXArray`
///
/// ### See Also
/// - <doc:lazy-evaluation>
public func eval(_ arrays: MLXArray...) {
let vector_array = new_mlx_vector_array(arrays)
_ = withEvalLock {
let result = withEvalLock {
mlx_async_eval(vector_array)
}
if result == 0 {
mlx_eval(vector_array)
}
mlx_vector_array_free(vector_array)
Expand All @@ -133,7 +151,10 @@ public func eval(_ arrays: MLXArray...) {
/// - <doc:lazy-evaluation>
public func eval(_ arrays: some Collection<MLXArray>) {
let vector_array = new_mlx_vector_array(arrays)
_ = withEvalLock {
let result = withEvalLock {
mlx_async_eval(vector_array)
}
if result == 0 {
mlx_eval(vector_array)
}
mlx_vector_array_free(vector_array)
Expand Down
Loading
Loading