Fix race condition during the tracing phase of automatic differentiation (vjp, value_and_grad) (#338)

* Fix race condition during the tracing phase of automatic differentiation (vjp, value_and_grad)

While the execution of the graph is thread-safe and lazy, the construction of that graph is not.

Author: aleroot

Source/MLX/Transforms+Internal.swift

@@ -15,7 +15,10 @@ private func valueAndGradient(
 
     var r0 = mlx_vector_array_new()
     var r1 = mlx_vector_array_new()
-    mlx_closure_value_and_grad_apply(&r0, &r1, valueAndGrad, input_vector)
+
+    _ = evalLock.withLock {
+        mlx_closure_value_and_grad_apply(&r0, &r1, valueAndGrad, input_vector)
+    }
 
     defer { mlx_vector_array_free(r0) }
     defer { mlx_vector_array_free(r1) }

Source/MLX/Transforms.swift

@@ -27,7 +27,11 @@ public func jvp(
     var r1 = mlx_vector_array_new()
 
     let closure = new_mlx_closure(f)
-    mlx_jvp(&r0, &r1, closure, primals_mlx, tangents_mlx)
+
+    _ = evalLock.withLock {
+        mlx_jvp(&r0, &r1, closure, primals_mlx, tangents_mlx)
+    }
+
     mlx_closure_free(closure)
 
     defer { mlx_vector_array_free(r0) }
@@ -60,7 +64,11 @@ public func vjp(
     var r1 = mlx_vector_array_new()
 
     let closure = new_mlx_closure(f)
-    mlx_vjp(&r0, &r1, closure, primals_mlx, cotangents_mlx)
+
+    _ = evalLock.withLock {
+        mlx_vjp(&r0, &r1, closure, primals_mlx, cotangents_mlx)
+    }
+
     mlx_closure_free(closure)
 
     defer { mlx_vector_array_free(r0) }