diff --git a/Source/MLX/Cmlx+Util.swift b/Source/MLX/Cmlx+Util.swift
index 9279a74f..ae9c353c 100644
--- a/Source/MLX/Cmlx+Util.swift
+++ b/Source/MLX/Cmlx+Util.swift
@@ -10,6 +10,8 @@ func new_mlx_vector_array(_ arrays: some Collection<MLXArray>) -> mlx_vector_arr
     }
 }
 
+/// Convert `mlx_vector_array` into a Swift array of `MLXArray`.
+/// - Parameter vector_array: input vector
 func mlx_vector_array_values(_ vector_array: mlx_vector_array) -> [MLXArray] {
     (0 ..< mlx_vector_array_size(vector_array))
         .map { index in
@@ -20,6 +22,27 @@ func mlx_vector_array_values(_ vector_array: mlx_vector_array) -> [MLXArray] {
         }
 }
 
+/// Convert `mlx_vector_array` into a Swift array of `MLXArray` and ensure
+/// that there are at least `count` items in the result.
+///
+/// This variant is useful for callers that want to assume there are at
+/// least `count` values in the result, even though an error condition might
+/// result in fewer.
+///
+/// Note: the extra values supplied after the actual length of the vector
+/// are `mlxNone` and are not typically usable.
+///
+/// - Parameters:
+///   - vector_array: input vector
+///   - count: minimum number of values in the result
+func mlx_vector_array_values(
+    _ vector_array: mlx_vector_array, count: Int
+) -> [MLXArray] {
+    var result = mlx_vector_array_values(vector_array)
+    while result.count < count { result.append(.mlxNone) }
+    return result
+}
+
 func mlx_map_array_values(_ mlx_map: mlx_map_string_to_array) -> [String: MLXArray] {
     var result = [String: MLXArray]()
 
diff --git a/Source/MLX/Linalg.swift b/Source/MLX/Linalg.swift
index 3acc4571..c7e80723 100644
--- a/Source/MLX/Linalg.swift
+++ b/Source/MLX/Linalg.swift
@@ -223,7 +223,7 @@ public enum MLXLinalg {
         mlx_linalg_svd(&vec, array.ctx, true, stream.ctx)
         defer { mlx_vector_array_free(vec) }
 
-        let arrays = mlx_vector_array_values(vec)
+        let arrays = mlx_vector_array_values(vec, count: 3)
         return (arrays[0], arrays[1], arrays[2])
     }
 
@@ -242,7 +242,7 @@ public enum MLXLinalg {
         mlx_linalg_svd(&vec, array.ctx, false, stream.ctx)
         defer { mlx_vector_array_free(vec) }
 
-        let arrays = mlx_vector_array_values(vec)
+        let arrays = mlx_vector_array_values(vec, count: 1)
         return arrays[0]
     }
 
@@ -369,7 +369,7 @@ public enum MLXLinalg {
         var vec = mlx_vector_array_new()
         mlx_linalg_lu(&vec, a.ctx, stream.ctx)
         defer { mlx_vector_array_free(vec) }
-        let arrays = mlx_vector_array_values(vec)
+        let arrays = mlx_vector_array_values(vec, count: 3)
         return (arrays[0], arrays[1], arrays[2])
     }
 
diff --git a/Source/MLX/MLXArray+Indexing.swift b/Source/MLX/MLXArray+Indexing.swift
index 394ab40d..06f3972d 100644
--- a/Source/MLX/MLXArray+Indexing.swift
+++ b/Source/MLX/MLXArray+Indexing.swift
@@ -3,6 +3,8 @@
 import Cmlx
 import Foundation
 
+private let log = MLXLogger(label: "MLXArray+Indexing")
+
 private func int32Range(_ count: Int) -> [Int32] {
     Array(0 ..< Int32(count))
 }
@@ -334,6 +336,13 @@ extension MLXArray {
             // conform the shapes to each other
             arrayIndices = MLX.broadcast(arrays: arrayIndices)
 
+            guard !arrayIndices.isEmpty else {
+                // if we get to this point it means that there is an error
+                // handler that will catch the error
+                MLX.log.error("broadcast failed in subscript set")
+                return
+            }
+
             // the shape of the broadcast value
             let broadcastShape =
                 arrayIndices[0].shape.asInt32 + int32Range(axis + 1, self.ndim).map { dim($0) }
@@ -1154,7 +1163,16 @@ func scatterArguments(
     }
 
     // Broadcast the update to the indices and slices
-    arrayIndices = broadcast(arrays: arrayIndices)
+    arrayIndices = broadcast(arrays: arrayIndices, stream: stream)
+
+    guard !arrayIndices.isEmpty else {
+        // if we get to this point it means that there is an error
+        // handler that will catch the error
+        MLX.log.error("broadcast failed in scatterArguments")
+        // empty indices + src as the update -> caller does src._updateInternal(src), a no-op
+        return ([], src, [])
+    }
+
     let updateShapeBroadcast =
         arrayIndices[0].shape + sliceShapes + src.shape.dropFirst(nonNewAxisOperationCount)
     update = broadcast(update, to: updateShapeBroadcast, stream: stream)
diff --git a/Source/MLX/MLXArray+Ops.swift b/Source/MLX/MLXArray+Ops.swift
index bb80a16f..3c07610c 100644
--- a/Source/MLX/MLXArray+Ops.swift
+++ b/Source/MLX/MLXArray+Ops.swift
@@ -2767,7 +2767,7 @@ extension MLXArray {
         var vec = mlx_vector_array_new()
         mlx_split(&vec, ctx, 2, axis.int32, stream.ctx)
         defer { mlx_vector_array_free(vec) }
-        let pieces = mlx_vector_array_values(vec)
+        let pieces = mlx_vector_array_values(vec, count: 2)
         return (pieces[0], pieces[1])
     }
 
diff --git a/Source/MLX/MLXFast.swift b/Source/MLX/MLXFast.swift
index f4f697c7..ed634425 100644
--- a/Source/MLX/MLXFast.swift
+++ b/Source/MLX/MLXFast.swift
@@ -36,7 +36,7 @@ public enum MLXFast {
             (mask ?? .mlxNone).ctx,
             stream.ctx)
 
-        let tuple = mlx_vector_array_values(result)
+        let tuple = mlx_vector_array_values(result, count: 2)
         return (output: tuple[0], state: tuple[1])
     }
 
diff --git a/Source/MLX/Ops+Array.swift b/Source/MLX/Ops+Array.swift
index 98b48b7e..8520aeb7 100644
--- a/Source/MLX/Ops+Array.swift
+++ b/Source/MLX/Ops+Array.swift
@@ -1684,7 +1684,7 @@ public func split(_ array: MLXArray, axis: Int = 0, stream: StreamOrDevice = .de
     var vec = mlx_vector_array_new()
     mlx_split(&vec, array.ctx, 2, axis.int32, stream.ctx)
     defer { mlx_vector_array_free(vec) }
-    let pieces = mlx_vector_array_values(vec)
+    let pieces = mlx_vector_array_values(vec, count: 2)
     return (pieces[0], pieces[1])
 }
 
diff --git a/Source/MLX/Ops.swift b/Source/MLX/Ops.swift
index f653918a..1f4ace2b 100644
--- a/Source/MLX/Ops.swift
+++ b/Source/MLX/Ops.swift
@@ -1066,7 +1066,7 @@ public func depends(input: MLXArray, dependencies: [MLXArray]) -> MLXArray {
     defer { mlx_vector_array_free(dependencies) }
     mlx_depends(&result, inputs, dependencies)
 
-    let arrays = mlx_vector_array_values(result)
+    let arrays = mlx_vector_array_values(result, count: 1)
     return arrays[0]
 }
 
@@ -1235,7 +1235,7 @@ public func divmod(
     var vec = mlx_vector_array_new()
     mlx_divmod(&vec, a.ctx, b.ctx, stream.ctx)
     defer { mlx_vector_array_free(vec) }
-    let result = mlx_vector_array_values(vec)
+    let result = mlx_vector_array_values(vec, count: 2)
     return (result[0], result[1])
 }
 
@@ -2510,7 +2510,7 @@ public func quantized(
         (globalScale ?? .mlxNone).ctx,
         stream.ctx)
 
-    let arrays = mlx_vector_array_values(r)
+    let arrays = mlx_vector_array_values(r, count: 2)
     return (arrays[0], arrays[1], arrays.count > 2 ? arrays[2] : nil)
 }
 
@@ -3061,7 +3061,7 @@ public func uniqued(
         (fillValue ?? .mlxNone).ctx,
         stream.ctx)
 
-    let arrays = mlx_vector_array_values(vector)
+    let arrays = mlx_vector_array_values(vector, count: 1)
     return arrays[0]
 }
 
@@ -3159,7 +3159,9 @@ public func uniqued(
         (fillValue ?? .mlxNone).ctx,
         stream.ctx)
 
-    let arrays = mlx_vector_array_values(vector)
+    let expectedCount =
+        1 + [UniqueOptions.index, .inverse, .counts].filter { returning.contains($0) }.count
+    let arrays = mlx_vector_array_values(vector, count: expectedCount)
     var result = UniqueResult(values: arrays[0])
 
     var readIndex = 1
diff --git a/Source/MLX/Transforms+Compile.swift b/Source/MLX/Transforms+Compile.swift
index d9ce19d4..293e6ac3 100644
--- a/Source/MLX/Transforms+Compile.swift
+++ b/Source/MLX/Transforms+Compile.swift
@@ -47,7 +47,7 @@ final class CompiledFunction: @unchecked (Sendable) {
 
     /// guards the observed state (``inputs`` / ``outputs``) -- see #226
     ///
-    /// This is always acquired *inside* `evalLock`; see ``call(_:)`` for the
+    /// This is always acquired *inside* `evalLock`; see ``call(_:expected:)`` for the
     /// ordering rule.  It is recursive because a compiled function's body can
     /// legitimately re-enter that same compiled function on the same thread
     /// while it is being traced (a nested call with tracer inputs bypasses the
@@ -156,8 +156,15 @@ final class CompiledFunction: @unchecked (Sendable) {
     /// hold it through `mlx_closure_apply` -- i.e. across the actual evaluation
     /// -- so this only moves the acquisition earlier by the state gathering and
     /// closure creation at the top of ``innerCall(_:)``.
-    func call(_ arguments: [MLXArray]) -> [MLXArray] {
-        withEvalLock {
+    ///
+    /// - Parameters:
+    ///   - arguments: arguments to the compiled function
+    ///   - expected: minimum number of results.  If an MLX error fires inside a
+    ///     `withError` scope the error is stored in the ErrorBox and the results
+    ///     are padded with `mlxNone` so callers can index them without crashing
+    ///     and `withError` can throw.
+    func call(_ arguments: [MLXArray], expected: Int) -> [MLXArray] {
+        var results = withEvalLock {
             // Stateless functions need only evalLock for cache bookkeeping.
             if inputs.isEmpty && outputs.isEmpty {
                 return innerCall(arguments)
@@ -166,11 +173,13 @@ final class CompiledFunction: @unchecked (Sendable) {
                 innerCall(arguments)
             }
         }
+        while results.count < expected { results.append(.mlxNone) }
+        return results
     }
 
     /// Acquire ``lock``, checking the ordering invariant first.
     ///
-    /// The check is where the invariant actually lives, not in ``call(_:)``: a
+    /// The check is where the invariant actually lives, not in ``call(_:expected:)``: a
     /// change that moved the per-instance lock back outside `evalLock` -- the
     /// shape that produced the deadlock -- trips this on the first compiled
     /// call, deterministically and without needing to lose a race.
@@ -184,7 +193,7 @@ final class CompiledFunction: @unchecked (Sendable) {
     }
 
     /// - Precondition: `evalLock` is held by the caller, and ``lock`` is held
-    ///   too unless this function has no observed state; see ``call(_:)``,
+    ///   too unless this function has no observed state; see ``call(_:expected:)``,
     ///   which is the only caller.
     func innerCall(_ arguments: [MLXArray]) -> [MLXArray] {
         let cache = ThreadCompilationCache.current
@@ -260,7 +269,7 @@ final class CompiledFunction: @unchecked (Sendable) {
         // note: this will use the cached compile (via the id)
         // but will be able to re-evaluate with fresh state if needed
         //
-        // evalLock is already held by call(_:) and covers everything through
+        // evalLock is already held by call(_:expected:) and covers everything through
         // mlx_closure_apply below, which is where the trace runs.
         var compiled = mlx_closure_new()
         let compileStatus = mlx_detail_compile(&compiled, innerClosure, id, shapeless, [], 0)
@@ -308,6 +317,13 @@ final class CompiledFunction: @unchecked (Sendable) {
         // push the stateOutput into the state
         let stateOutput = outputs.flatMap { $0.innerState() }.filter { !$0.isSealed }
 
+        // a short result means something went wrong -- don't push result
+        // arrays into the state (and avoid a negative prefix below).
+        // call(_:expected:) pads the results.
+        guard resultsPlusStateOutput.count >= stateOutput.count else {
+            return []
+        }
+
         for (s, newValues) in zip(stateOutput, resultsPlusStateOutput.suffix(stateOutput.count)) {
             s._updateInternal(newValues)
         }
@@ -344,7 +360,7 @@ public func compile(
     let compileState = CompiledFunction(inputs: inputs, outputs: outputs, shapeless: shapeless, f)
 
     return { arrays in
-        compileState.call(arrays)
+        compileState.call(arrays, expected: 0)
     }
 }
 
@@ -369,7 +385,7 @@ public func compile(
     let compileState = CompiledFunction(shapeless: shapeless, f)
 
     return { arrays in
-        compileState.call(arrays)
+        compileState.call(arrays, expected: 0)
     }
 }
 
@@ -388,11 +404,7 @@ public func compile(
     }
 
     return { a in
-        let r = compileState.call([a])
-        // r is empty only when an MLX error fired inside a withError scope — the
-        // error is already stored in the ErrorBox.  Return a placeholder so that
-        // withError can throw instead of crashing with "Index out of range".
-        return r.isEmpty ? MLXArray(0) : r[0]
+        compileState.call([a], expected: 1)[0]
     }
 }
 
@@ -411,11 +423,7 @@ public func compile(
     }
 
     return { a in
-        let r = compileState.call([a])
-        // r is empty only when an MLX error fired inside a withError scope — the
-        // error is already stored in the ErrorBox.  Return a placeholder so that
-        // withError can throw instead of crashing with "Index out of range".
-        return r.isEmpty ? MLXArray(0) : r[0]
+        compileState.call([a], expected: 1)[0]
     }
 }
 
@@ -436,8 +444,7 @@ public func compile(
     }
 
     return { a, b in
-        let r = compileState.call([a, b])
-        return r.isEmpty ? MLXArray(0) : r[0]
+        compileState.call([a, b], expected: 1)[0]
     }
 }
 
@@ -458,8 +465,7 @@ public func compile(
     }
 
     return { a, b in
-        let r = compileState.call([a, b])
-        return r.isEmpty ? MLXArray(0) : r[0]
+        compileState.call([a, b], expected: 1)[0]
     }
 }
 
@@ -480,8 +486,7 @@ public func compile(
     }
 
     return { a, b, c in
-        let r = compileState.call([a, b, c])
-        return r.isEmpty ? MLXArray(0) : r[0]
+        compileState.call([a, b, c], expected: 1)[0]
     }
 }
 
@@ -502,8 +507,7 @@ public func compile(
     }
 
     return { a, b, c in
-        let r = compileState.call([a, b, c])
-        return r.isEmpty ? MLXArray(0) : r[0]
+        compileState.call([a, b, c], expected: 1)[0]
     }
 }
 
diff --git a/Source/MLX/Transforms+CompileOverloads.swift b/Source/MLX/Transforms+CompileOverloads.swift
index 2f478dfb..619c9821 100644
--- a/Source/MLX/Transforms+CompileOverloads.swift
+++ b/Source/MLX/Transforms+CompileOverloads.swift
@@ -15,7 +15,7 @@ public func compile(
         return [r.0, r.1]
     }
     return { a in
-        let r = compileState.call([a])
+        let r = compileState.call([a], expected: 2)
         return (r[0], r[1])
     }
 }
@@ -31,7 +31,7 @@ public func compile(
         return [r.0, r.1, r.2]
     }
     return { a in
-        let r = compileState.call([a])
+        let r = compileState.call([a], expected: 3)
         return (r[0], r[1], r[2])
     }
 }
@@ -47,7 +47,7 @@ public func compile(
         return [r.0, r.1, r.2, r.3]
     }
     return { a in
-        let r = compileState.call([a])
+        let r = compileState.call([a], expected: 4)
         return (r[0], r[1], r[2], r[3])
     }
 }
@@ -63,7 +63,7 @@ public func compile(
         return [r.0, r.1]
     }
     return { a, b in
-        let r = compileState.call([a, b])
+        let r = compileState.call([a, b], expected: 2)
         return (r[0], r[1])
     }
 }
@@ -79,7 +79,7 @@ public func compile(
         return [r.0, r.1, r.2]
     }
     return { a, b in
-        let r = compileState.call([a, b])
+        let r = compileState.call([a, b], expected: 3)
         return (r[0], r[1], r[2])
     }
 }
@@ -95,7 +95,7 @@ public func compile(
         return [r.0, r.1, r.2, r.3]
     }
     return { a, b in
-        let r = compileState.call([a, b])
+        let r = compileState.call([a, b], expected: 4)
         return (r[0], r[1], r[2], r[3])
     }
 }
@@ -111,7 +111,7 @@ public func compile(
         return [r.0, r.1]
     }
     return { a, b, c in
-        let r = compileState.call([a, b, c])
+        let r = compileState.call([a, b, c], expected: 2)
         return (r[0], r[1])
     }
 }
@@ -127,7 +127,7 @@ public func compile(
         return [r.0, r.1, r.2]
     }
     return { a, b, c in
-        let r = compileState.call([a, b, c])
+        let r = compileState.call([a, b, c], expected: 3)
         return (r[0], r[1], r[2])
     }
 }
@@ -146,7 +146,7 @@ public func compile(
         return [r.0, r.1, r.2, r.3]
     }
     return { a, b, c in
-        let r = compileState.call([a, b, c])
+        let r = compileState.call([a, b, c], expected: 4)
         return (r[0], r[1], r[2], r[3])
     }
 }
@@ -161,7 +161,7 @@ public func compile(
         [f(args[0], args[1], args[2], args[3])]
     }
     return { a, b, c, d in
-        compileState.call([a, b, c, d])[0]
+        compileState.call([a, b, c, d], expected: 1)[0]
     }
 }
 
@@ -176,7 +176,7 @@ public func compile(
         return [r.0, r.1]
     }
     return { a, b, c, d in
-        let r = compileState.call([a, b, c, d])
+        let r = compileState.call([a, b, c, d], expected: 2)
         return (r[0], r[1])
     }
 }
@@ -195,7 +195,7 @@ public func compile(
         return [r.0, r.1, r.2]
     }
     return { a, b, c, d in
-        let r = compileState.call([a, b, c, d])
+        let r = compileState.call([a, b, c, d], expected: 3)
         return (r[0], r[1], r[2])
     }
 }
@@ -214,7 +214,7 @@ public func compile(
         return [r.0, r.1, r.2, r.3]
     }
     return { a, b, c, d in
-        let r = compileState.call([a, b, c, d])
+        let r = compileState.call([a, b, c, d], expected: 4)
         return (r[0], r[1], r[2], r[3])
     }
 }
@@ -229,7 +229,7 @@ public func compile(
         [f(args[0], args[1], args[2], args[3], args[4])]
     }
     return { a, b, c, d, e in
-        compileState.call([a, b, c, d, e])[0]
+        compileState.call([a, b, c, d, e], expected: 1)[0]
     }
 }
 
@@ -247,7 +247,7 @@ public func compile(
         return [r.0, r.1]
     }
     return { a, b, c, d, e in
-        let r = compileState.call([a, b, c, d, e])
+        let r = compileState.call([a, b, c, d, e], expected: 2)
         return (r[0], r[1])
     }
 }
@@ -266,7 +266,7 @@ public func compile(
         return [r.0, r.1, r.2]
     }
     return { a, b, c, d, e in
-        let r = compileState.call([a, b, c, d, e])
+        let r = compileState.call([a, b, c, d, e], expected: 3)
         return (r[0], r[1], r[2])
     }
 }
@@ -289,7 +289,7 @@ public func compile(
         return [r.0, r.1, r.2, r.3]
     }
     return { a, b, c, d, e in
-        let r = compileState.call([a, b, c, d, e])
+        let r = compileState.call([a, b, c, d, e], expected: 4)
         return (r[0], r[1], r[2], r[3])
     }
 }
@@ -305,7 +305,7 @@ public func compile(
         [f(args[0], args[1], args[2], args[3], args[4], args[5])]
     }
     return { a, b, c, d, e, g in
-        compileState.call([a, b, c, d, e, g])[0]
+        compileState.call([a, b, c, d, e, g], expected: 1)[0]
     }
 }
 
@@ -323,7 +323,7 @@ public func compile(
         return [r.0, r.1]
     }
     return { a, b, c, d, e, g in
-        let r = compileState.call([a, b, c, d, e, g])
+        let r = compileState.call([a, b, c, d, e, g], expected: 2)
         return (r[0], r[1])
     }
 }
@@ -346,7 +346,7 @@ public func compile(
         return [r.0, r.1, r.2]
     }
     return { a, b, c, d, e, g in
-        let r = compileState.call([a, b, c, d, e, g])
+        let r = compileState.call([a, b, c, d, e, g], expected: 3)
         return (r[0], r[1], r[2])
     }
 }
@@ -369,7 +369,7 @@ public func compile(
         return [r.0, r.1, r.2, r.3]
     }
     return { a, b, c, d, e, g in
-        let r = compileState.call([a, b, c, d, e, g])
+        let r = compileState.call([a, b, c, d, e, g], expected: 4)
         return (r[0], r[1], r[2], r[3])
     }
 }
@@ -386,7 +386,7 @@ public func compile(
         [f(args[0], args[1], args[2], args[3], args[4], args[5], args[6])]
     }
     return { a, b, c, d, e, g, h in
-        compileState.call([a, b, c, d, e, g, h])[0]
+        compileState.call([a, b, c, d, e, g, h], expected: 1)[0]
     }
 }
 
@@ -409,7 +409,7 @@ public func compile(
         return [r.0, r.1]
     }
     return { a, b, c, d, e, g, h in
-        let r = compileState.call([a, b, c, d, e, g, h])
+        let r = compileState.call([a, b, c, d, e, g, h], expected: 2)
         return (r[0], r[1])
     }
 }
@@ -433,7 +433,7 @@ public func compile(
         return [r.0, r.1, r.2]
     }
     return { a, b, c, d, e, g, h in
-        let r = compileState.call([a, b, c, d, e, g, h])
+        let r = compileState.call([a, b, c, d, e, g, h], expected: 3)
         return (r[0], r[1], r[2])
     }
 }
@@ -457,7 +457,7 @@ public func compile(
         return [r.0, r.1, r.2, r.3]
     }
     return { a, b, c, d, e, g, h in
-        let r = compileState.call([a, b, c, d, e, g, h])
+        let r = compileState.call([a, b, c, d, e, g, h], expected: 4)
         return (r[0], r[1], r[2], r[3])
     }
 }
@@ -479,7 +479,7 @@ public func compile(
         [f(args[0], args[1], args[2], args[3], args[4], args[5], args[6], args[7])]
     }
     return { a, b, c, d, e, g, h, i in
-        compileState.call([a, b, c, d, e, g, h, i])[0]
+        compileState.call([a, b, c, d, e, g, h, i], expected: 1)[0]
     }
 }
 
@@ -501,7 +501,7 @@ public func compile(
         return [r.0, r.1]
     }
     return { a, b, c, d, e, g, h, i in
-        let r = compileState.call([a, b, c, d, e, g, h, i])
+        let r = compileState.call([a, b, c, d, e, g, h, i], expected: 2)
         return (r[0], r[1])
     }
 }
@@ -524,7 +524,7 @@ public func compile(
         return [r.0, r.1, r.2]
     }
     return { a, b, c, d, e, g, h, i in
-        let r = compileState.call([a, b, c, d, e, g, h, i])
+        let r = compileState.call([a, b, c, d, e, g, h, i], expected: 3)
         return (r[0], r[1], r[2])
     }
 }
@@ -547,7 +547,7 @@ public func compile(
         return [r.0, r.1, r.2, r.3]
     }
     return { a, b, c, d, e, g, h, i in
-        let r = compileState.call([a, b, c, d, e, g, h, i])
+        let r = compileState.call([a, b, c, d, e, g, h, i], expected: 4)
         return (r[0], r[1], r[2], r[3])
     }
 }
@@ -565,7 +565,7 @@ public func compile(
         return [r.0, r.1]
     }
     return { a in
-        let r = compileState.call([a])
+        let r = compileState.call([a], expected: 2)
         return (r[0], r[1])
     }
 }
@@ -581,7 +581,7 @@ public func compile(
         return [r.0, r.1, r.2]
     }
     return { a in
-        let r = compileState.call([a])
+        let r = compileState.call([a], expected: 3)
         return (r[0], r[1], r[2])
     }
 }
@@ -597,7 +597,7 @@ public func compile(
         return [r.0, r.1, r.2, r.3]
     }
     return { a in
-        let r = compileState.call([a])
+        let r = compileState.call([a], expected: 4)
         return (r[0], r[1], r[2], r[3])
     }
 }
@@ -613,7 +613,7 @@ public func compile(
         return [r.0, r.1]
     }
     return { a, b in
-        let r = compileState.call([a, b])
+        let r = compileState.call([a, b], expected: 2)
         return (r[0], r[1])
     }
 }
@@ -629,7 +629,7 @@ public func compile(
         return [r.0, r.1, r.2]
     }
     return { a, b in
-        let r = compileState.call([a, b])
+        let r = compileState.call([a, b], expected: 3)
         return (r[0], r[1], r[2])
     }
 }
@@ -645,7 +645,7 @@ public func compile(
         return [r.0, r.1, r.2, r.3]
     }
     return { a, b in
-        let r = compileState.call([a, b])
+        let r = compileState.call([a, b], expected: 4)
         return (r[0], r[1], r[2], r[3])
     }
 }
@@ -661,7 +661,7 @@ public func compile(
         return [r.0, r.1]
     }
     return { a, b, c in
-        let r = compileState.call([a, b, c])
+        let r = compileState.call([a, b, c], expected: 2)
         return (r[0], r[1])
     }
 }
@@ -677,7 +677,7 @@ public func compile(
         return [r.0, r.1, r.2]
     }
     return { a, b, c in
-        let r = compileState.call([a, b, c])
+        let r = compileState.call([a, b, c], expected: 3)
         return (r[0], r[1], r[2])
     }
 }
@@ -696,7 +696,7 @@ public func compile(
         return [r.0, r.1, r.2, r.3]
     }
     return { a, b, c in
-        let r = compileState.call([a, b, c])
+        let r = compileState.call([a, b, c], expected: 4)
         return (r[0], r[1], r[2], r[3])
     }
 }
@@ -711,7 +711,7 @@ public func compile(
         [f(args[0], args[1], args[2], args[3])]
     }
     return { a, b, c, d in
-        compileState.call([a, b, c, d])[0]
+        compileState.call([a, b, c, d], expected: 1)[0]
     }
 }
 
@@ -726,7 +726,7 @@ public func compile(
         return [r.0, r.1]
     }
     return { a, b, c, d in
-        let r = compileState.call([a, b, c, d])
+        let r = compileState.call([a, b, c, d], expected: 2)
         return (r[0], r[1])
     }
 }
@@ -745,7 +745,7 @@ public func compile(
         return [r.0, r.1, r.2]
     }
     return { a, b, c, d in
-        let r = compileState.call([a, b, c, d])
+        let r = compileState.call([a, b, c, d], expected: 3)
         return (r[0], r[1], r[2])
     }
 }
@@ -765,7 +765,7 @@ public func compile(
         return [r.0, r.1, r.2, r.3]
     }
     return { a, b, c, d in
-        let r = compileState.call([a, b, c, d])
+        let r = compileState.call([a, b, c, d], expected: 4)
         return (r[0], r[1], r[2], r[3])
     }
 }
@@ -780,7 +780,7 @@ public func compile(
         [f(args[0], args[1], args[2], args[3], args[4])]
     }
     return { a, b, c, d, e in
-        compileState.call([a, b, c, d, e])[0]
+        compileState.call([a, b, c, d, e], expected: 1)[0]
     }
 }
 
@@ -798,7 +798,7 @@ public func compile(
         return [r.0, r.1]
     }
     return { a, b, c, d, e in
-        let r = compileState.call([a, b, c, d, e])
+        let r = compileState.call([a, b, c, d, e], expected: 2)
         return (r[0], r[1])
     }
 }
@@ -818,7 +818,7 @@ public func compile(
         return [r.0, r.1, r.2]
     }
     return { a, b, c, d, e in
-        let r = compileState.call([a, b, c, d, e])
+        let r = compileState.call([a, b, c, d, e], expected: 3)
         return (r[0], r[1], r[2])
     }
 }
@@ -841,7 +841,7 @@ public func compile(
         return [r.0, r.1, r.2, r.3]
     }
     return { a, b, c, d, e in
-        let r = compileState.call([a, b, c, d, e])
+        let r = compileState.call([a, b, c, d, e], expected: 4)
         return (r[0], r[1], r[2], r[3])
     }
 }
@@ -857,7 +857,7 @@ public func compile(
         [f(args[0], args[1], args[2], args[3], args[4], args[5])]
     }
     return { a, b, c, d, e, g in
-        compileState.call([a, b, c, d, e, g])[0]
+        compileState.call([a, b, c, d, e, g], expected: 1)[0]
     }
 }
 
@@ -876,7 +876,7 @@ public func compile(
         return [r.0, r.1]
     }
     return { a, b, c, d, e, g in
-        let r = compileState.call([a, b, c, d, e, g])
+        let r = compileState.call([a, b, c, d, e, g], expected: 2)
         return (r[0], r[1])
     }
 }
@@ -899,7 +899,7 @@ public func compile(
         return [r.0, r.1, r.2]
     }
     return { a, b, c, d, e, g in
-        let r = compileState.call([a, b, c, d, e, g])
+        let r = compileState.call([a, b, c, d, e, g], expected: 3)
         return (r[0], r[1], r[2])
     }
 }
@@ -922,7 +922,7 @@ public func compile(
         return [r.0, r.1, r.2, r.3]
     }
     return { a, b, c, d, e, g in
-        let r = compileState.call([a, b, c, d, e, g])
+        let r = compileState.call([a, b, c, d, e, g], expected: 4)
         return (r[0], r[1], r[2], r[3])
     }
 }
@@ -939,7 +939,7 @@ public func compile(
         [f(args[0], args[1], args[2], args[3], args[4], args[5], args[6])]
     }
     return { a, b, c, d, e, g, h in
-        compileState.call([a, b, c, d, e, g, h])[0]
+        compileState.call([a, b, c, d, e, g, h], expected: 1)[0]
     }
 }
 
@@ -962,7 +962,7 @@ public func compile(
         return [r.0, r.1]
     }
     return { a, b, c, d, e, g, h in
-        let r = compileState.call([a, b, c, d, e, g, h])
+        let r = compileState.call([a, b, c, d, e, g, h], expected: 2)
         return (r[0], r[1])
     }
 }
@@ -986,7 +986,7 @@ public func compile(
         return [r.0, r.1, r.2]
     }
     return { a, b, c, d, e, g, h in
-        let r = compileState.call([a, b, c, d, e, g, h])
+        let r = compileState.call([a, b, c, d, e, g, h], expected: 3)
         return (r[0], r[1], r[2])
     }
 }
@@ -1010,7 +1010,7 @@ public func compile(
         return [r.0, r.1, r.2, r.3]
     }
     return { a, b, c, d, e, g, h in
-        let r = compileState.call([a, b, c, d, e, g, h])
+        let r = compileState.call([a, b, c, d, e, g, h], expected: 4)
         return (r[0], r[1], r[2], r[3])
     }
 }
@@ -1032,7 +1032,7 @@ public func compile(
         [f(args[0], args[1], args[2], args[3], args[4], args[5], args[6], args[7])]
     }
     return { a, b, c, d, e, g, h, i in
-        compileState.call([a, b, c, d, e, g, h, i])[0]
+        compileState.call([a, b, c, d, e, g, h, i], expected: 1)[0]
     }
 }
 
@@ -1054,7 +1054,7 @@ public func compile(
         return [r.0, r.1]
     }
     return { a, b, c, d, e, g, h, i in
-        let r = compileState.call([a, b, c, d, e, g, h, i])
+        let r = compileState.call([a, b, c, d, e, g, h, i], expected: 2)
         return (r[0], r[1])
     }
 }
@@ -1077,7 +1077,7 @@ public func compile(
         return [r.0, r.1, r.2]
     }
     return { a, b, c, d, e, g, h, i in
-        let r = compileState.call([a, b, c, d, e, g, h, i])
+        let r = compileState.call([a, b, c, d, e, g, h, i], expected: 3)
         return (r[0], r[1], r[2])
     }
 }
@@ -1100,7 +1100,7 @@ public func compile(
         return [r.0, r.1, r.2, r.3]
     }
     return { a, b, c, d, e, g, h, i in
-        let r = compileState.call([a, b, c, d, e, g, h, i])
+        let r = compileState.call([a, b, c, d, e, g, h, i], expected: 4)
         return (r[0], r[1], r[2], r[3])
     }
 }
diff --git a/Source/MLX/Transforms+Grad.swift b/Source/MLX/Transforms+Grad.swift
index af1aecd8..0d41fe02 100644
--- a/Source/MLX/Transforms+Grad.swift
+++ b/Source/MLX/Transforms+Grad.swift
@@ -59,7 +59,6 @@ public func valueAndGrad(
 )
     -> ([MLXArray]) -> ([MLXArray], [MLXArray])
 {
-    //
     return buildValueAndGradient(f, argumentNumbers: argumentNumbers)
 }
 
diff --git a/Source/MLX/Transforms+Internal.swift b/Source/MLX/Transforms+Internal.swift
index c73ecf18..de2a64fb 100644
--- a/Source/MLX/Transforms+Internal.swift
+++ b/Source/MLX/Transforms+Internal.swift
@@ -5,11 +5,16 @@ import Foundation
 
 // see Transforms+Variants for generated grad() functions
 
+/// Inner build of the value and gradient function
+/// - Parameters:
+///   - valueAndGrad: closure to transform
+///   - arrays: inputs
+///   - expected: minimum number of values and gradients to return.  If an error
+///     occurs, the results are padded with `mlxNone` so callers can index safely.
 private func valueAndGradient(
-    apply valueAndGrad: mlx_closure_value_and_grad, arrays: some Collection<MLXArray>
-)
-    -> ([MLXArray], [MLXArray])
-{
+    apply valueAndGrad: mlx_closure_value_and_grad, arrays: some Collection<MLXArray>,
+    expected: (values: Int, gradients: Int)
+) -> ([MLXArray], [MLXArray]) {
     let input_vector = new_mlx_vector_array(arrays)
     defer { mlx_vector_array_free(input_vector) }
 
@@ -23,7 +28,10 @@ private func valueAndGradient(
     defer { mlx_vector_array_free(r0) }
     defer { mlx_vector_array_free(r1) }
 
-    return (mlx_vector_array_values(r0), mlx_vector_array_values(r1))
+    return (
+        mlx_vector_array_values(r0, count: expected.values),
+        mlx_vector_array_values(r1, count: expected.gradients)
+    )
 }
 
 func buildGradient(_ f: @escaping ([MLXArray]) -> [MLXArray], argumentNumbers: some Collection<Int>)
@@ -40,7 +48,11 @@ func buildGradient(_ f: @escaping ([MLXArray]) -> [MLXArray], argumentNumbers: s
 
         defer { mlx_closure_value_and_grad_free(vag) }
 
-        return valueAndGradient(apply: vag, arrays: arrays).1
+        // unwrapResult() reads gradients[0], so ensure at least one
+        let (_, gradients) = valueAndGradient(
+            apply: vag, arrays: arrays,
+            expected: (values: 0, gradients: max(1, argumentNumbers.count)))
+        return gradients
     }
 }
 
@@ -58,7 +70,9 @@ func buildValueAndGradient(
 
         defer { mlx_closure_value_and_grad_free(vag) }
 
-        return valueAndGradient(apply: vag, arrays: arrays)
+        return valueAndGradient(
+            apply: vag, arrays: arrays,
+            expected: (values: 1, gradients: argumentNumbers.count))
     }
 }
 
@@ -106,7 +120,9 @@ func buildValueAndGradient<T>(
 
         defer { mlx_closure_value_and_grad_free(vag) }
 
-        let (values, flatGradients) = valueAndGradient(apply: vag, arrays: flattenedArrays)
+        let (values, flatGradients) = valueAndGradient(
+            apply: vag, arrays: flattenedArrays,
+            expected: (values: 1, gradients: flattenedArrays.count))
         let gradients = unflattened(flatGradients)
 
         return (values, gradients)
diff --git a/Source/MLX/Transforms.swift b/Source/MLX/Transforms.swift
index e9e8c7b6..9deed823 100644
--- a/Source/MLX/Transforms.swift
+++ b/Source/MLX/Transforms.swift
@@ -74,7 +74,11 @@ public func vjp(
     defer { mlx_vector_array_free(r0) }
     defer { mlx_vector_array_free(r1) }
 
-    return (mlx_vector_array_values(r0), mlx_vector_array_values(r1))
+    // outputs match the cotangents in number, vjps match the primals
+    return (
+        mlx_vector_array_values(r0, count: cotangents.count),
+        mlx_vector_array_values(r1, count: primals.count)
+    )
 }
 
 /// Returns a function that computes the gradient and result of `f`, computing the gradient with respect to the ``NestedDictionary``.
diff --git a/Tests/MLXTests/ArrayAtTests.swift b/Tests/MLXTests/ArrayAtTests.swift
index 57b077e9..740a031a 100644
--- a/Tests/MLXTests/ArrayAtTests.swift
+++ b/Tests/MLXTests/ArrayAtTests.swift
@@ -2,11 +2,12 @@
 
 import Foundation
 import MLX
-import XCTest
+import Testing
 
-class ArrayAtTests: XCTestCase {
+@Suite
+struct ArrayAtTests {
 
-    func testArrayAt() {
+    @Test func arrayAt() {
         // from example at https://ml-explore.github.io/mlx/build/html/python/_autosummary/mlx.core.array.at.html#mlx.core.array.at
 
         // this references each index twice
@@ -15,27 +16,27 @@ class ArrayAtTests: XCTestCase {
         // assign through index -- we can only observe the last assignment to a location
         let a0 = MLXArray([0, 0])
         a0[idx] = MLXArray(2)
-        assertEqual(a0, MLXArray([2, 2]))
+        expectEqual(a0, MLXArray([2, 2]))
 
         // similar to above -- we can only observe one assignment, so we just get a +1
         // note: there was a bug in the += operator where the lhs was not inout and
         // this was producing [0, 0]
         let a1 = MLXArray([0, 0])
         a1[idx] += 1
-        assertEqual(a1, MLXArray([1, 1]))
+        expectEqual(a1, MLXArray([1, 1]))
 
         // the bare add produces a value for each index including the duplicates
         let a2 = MLXArray([0, 0])
-        assertEqual(a2[idx] + 1, MLXArray([1, 1, 1, 1]))
+        expectEqual(a2[idx] + 1, MLXArray([1, 1, 1, 1]))
 
         // but the assign back through the index will collapse the values down
         // into the same location -- this is the same as a2[idx] += 1
         a2[idx] = a2[idx] + 1
-        assertEqual(a2, MLXArray([1, 1]))
+        expectEqual(a2, MLXArray([1, 1]))
 
         // this will update 0 and 1 twice
         let a3 = MLXArray([0, 0])
-        assertEqual(a3.at[idx].add(1), MLXArray([2, 2]))
+        expectEqual(a3.at[idx].add(1), MLXArray([2, 2]))
     }
 
     // MARK: - slice updates
@@ -43,52 +44,52 @@ class ArrayAtTests: XCTestCase {
     // pure slice indices route through the mlx_slice_update_* family rather than a
     // scatter -- these check that path produces the same answers.
 
-    func testSliceAdd() {
+    @Test func sliceAdd() {
         let a = MLXArray(0 ..< 6)
-        assertEqual(a.at[1 ..< 4].add(10), MLXArray([0, 11, 12, 13, 4, 5]))
+        expectEqual(a.at[1 ..< 4].add(10), MLXArray([0, 11, 12, 13, 4, 5]))
     }
 
-    func testSliceSubtract() {
+    @Test func sliceSubtract() {
         let a = MLXArray(0 ..< 6)
-        assertEqual(a.at[1 ..< 4].subtract(1), MLXArray([0, 0, 1, 2, 4, 5]))
+        expectEqual(a.at[1 ..< 4].subtract(1), MLXArray([0, 0, 1, 2, 4, 5]))
     }
 
-    func testSliceMultiply() {
+    @Test func sliceMultiply() {
         let a = MLXArray(1 ..< 7)
-        assertEqual(a.at[2 ..< 5].multiply(10), MLXArray([1, 2, 30, 40, 50, 6]))
+        expectEqual(a.at[2 ..< 5].multiply(10), MLXArray([1, 2, 30, 40, 50, 6]))
     }
 
-    func testSliceDivide() {
+    @Test func sliceDivide() {
         let a = MLXArray(converting: [2.0, 4.0, 6.0, 8.0])
-        assertEqual(
+        expectEqual(
             a.at[1 ..< 3].divide(2), MLXArray(converting: [2.0, 2.0, 3.0, 8.0]), atol: 1e-6)
     }
 
-    func testSliceMinimumMaximum() {
+    @Test func sliceMinimumMaximum() {
         let a = MLXArray([5, 5, 5, 5])
-        assertEqual(a.at[0 ..< 2].minimum(3), MLXArray([3, 3, 5, 5]))
-        assertEqual(a.at[0 ..< 2].maximum(7), MLXArray([7, 7, 5, 5]))
+        expectEqual(a.at[0 ..< 2].minimum(3), MLXArray([3, 3, 5, 5]))
+        expectEqual(a.at[0 ..< 2].maximum(7), MLXArray([7, 7, 5, 5]))
     }
 
-    func testSliceIndexAdd() {
+    @Test func sliceIndexAdd() {
         // a single integer index is also expressible as a slice update
         let a = MLXArray(0 ..< 4)
-        assertEqual(a.at[1].add(100), MLXArray([0, 101, 2, 3]))
+        expectEqual(a.at[1].add(100), MLXArray([0, 101, 2, 3]))
     }
 
-    func testMultiDimensionalSliceAdd() {
+    @Test func multiDimensionalSliceAdd() {
         let a = MLXArray(0 ..< 6, [2, 3])
-        assertEqual(
+        expectEqual(
             a.at[0 ..< 1, 1 ..< 3].add(10),
             MLXArray([0, 11, 12, 3, 4, 5], [2, 3]))
     }
 
-    func testStridedSliceAdd() {
+    @Test func stridedSliceAdd() {
         let a = MLXArray(0 ..< 6)
-        assertEqual(a.at[.stride(by: 2)].add(10), MLXArray([10, 1, 12, 3, 14, 5]))
+        expectEqual(a.at[.stride(by: 2)].add(10), MLXArray([10, 1, 12, 3, 14, 5]))
     }
 
-    func testSliceAndScatterAgree() {
+    @Test func sliceAndScatterAgree() {
         // the slice path and the scatter path should produce the same result --
         // indexing with an array forces the scatter fallback
         let a = MLXArray(0 ..< 6)
@@ -96,6 +97,6 @@ class ArrayAtTests: XCTestCase {
         let viaSlice = a.at[1 ..< 4].add(10)
         let viaScatter = a.at[MLXArray([1, 2, 3])].add(10)
 
-        assertEqual(viaSlice, viaScatter)
+        expectEqual(viaSlice, viaScatter)
     }
 }
diff --git a/Tests/MLXTests/DTypeTests.swift b/Tests/MLXTests/DTypeTests.swift
index 5d50306c..5b4e446d 100644
--- a/Tests/MLXTests/DTypeTests.swift
+++ b/Tests/MLXTests/DTypeTests.swift
@@ -2,51 +2,52 @@
 
 import Foundation
 import MLX
-import XCTest
+import Testing
 
-class DTypeTests: XCTestCase {
+@Suite
+struct DTypeTests {
 
-    func testScalarAsMLXArrayInt() {
+    @Test func scalarAsMLXArrayInt() {
         // Int (which is really a 64 bit value) produces .int32 to match
         // python mlx
         let i = 10
         let intArray = i.asMLXArray(dtype: nil)
-        XCTAssertEqual(intArray.dtype, .int32)
+        #expect(intArray.dtype == .int32)
 
         // but using Int64 directly will give .int64
         let i64 = Int64(10)
         let int64Array = i64.asMLXArray(dtype: nil)
-        XCTAssertEqual(int64Array.dtype, .int64)
+        #expect(int64Array.dtype == .int64)
     }
 
-    func testScalarAsMLXArrayPromotionInt() {
+    @Test func scalarAsMLXArrayPromotionInt() {
         let i = 10
 
         // does not promote arrays to larger type,
         // e.g. ([1, 2, 3] dtype: .int8) + scalar(10)
         // won't promote the whole thing to int32.
-        XCTAssertEqual(i.asMLXArray(dtype: .int8).dtype, .int8)
-        XCTAssertEqual(i.asMLXArray(dtype: .uint32).dtype, .uint32)
-        XCTAssertEqual(i.asMLXArray(dtype: .int64).dtype, .int64)
+        #expect(i.asMLXArray(dtype: .int8).dtype == .int8)
+        #expect(i.asMLXArray(dtype: .uint32).dtype == .uint32)
+        #expect(i.asMLXArray(dtype: .int64).dtype == .int64)
 
-        XCTAssertEqual(i.asMLXArray(dtype: .bool).dtype, .int32)
+        #expect(i.asMLXArray(dtype: .bool).dtype == .int32)
 
-        XCTAssertEqual(i.asMLXArray(dtype: .float16).dtype, .float16)
+        #expect(i.asMLXArray(dtype: .float16).dtype == .float16)
     }
 
-    func testScalarAsMLXArrayPromotionFloat() {
+    @Test func scalarAsMLXArrayPromotionFloat() {
         let f: Float = 1.5
 
         // floats do promote to a matching float type
-        XCTAssertEqual(f.asMLXArray(dtype: .int8).dtype, .float32)
-        XCTAssertEqual(f.asMLXArray(dtype: .uint32).dtype, .float32)
-        XCTAssertEqual(f.asMLXArray(dtype: .int64).dtype, .float32)
+        #expect(f.asMLXArray(dtype: .int8).dtype == .float32)
+        #expect(f.asMLXArray(dtype: .uint32).dtype == .float32)
+        #expect(f.asMLXArray(dtype: .int64).dtype == .float32)
 
-        XCTAssertEqual(f.asMLXArray(dtype: .bool).dtype, .float32)
+        #expect(f.asMLXArray(dtype: .bool).dtype == .float32)
 
         // but ([1.5, 2.5, 3.5] dtype: .float16) + scalar(1.5)
         // doesn't promote the array to float32
-        XCTAssertEqual(f.asMLXArray(dtype: .float16).dtype, .float16)
+        #expect(f.asMLXArray(dtype: .float16).dtype == .float16)
     }
 
 }
diff --git a/Tests/MLXTests/DistributedNNTests.swift b/Tests/MLXTests/DistributedNNTests.swift
index 9d2ac861..523ef1ec 100644
--- a/Tests/MLXTests/DistributedNNTests.swift
+++ b/Tests/MLXTests/DistributedNNTests.swift
@@ -18,94 +18,95 @@ import XCTest
 func shardLinearBody(world: MLXDistributed.Group) throws {
     let rank = try world.rank
     let size = try world.size
-    MLXRandom.seed(0xF0F0_F0F0)
+    try withRandomState(MLXRandom.RandomState(seed: 0xF0F0_F0F0)) {
 
-    let lower = rank * 1024 / size
-    let upper = (rank + 1) * 1024 / size
+        let lower = rank * 1024 / size
+        let upper = (rank + 1) * 1024 / size
 
-    let x = MLXRandom.normal([4, 1024])
-    let linear = Linear(1024, 1024, bias: true)
+        let x = MLXRandom.normal([4, 1024])
+        let linear = Linear(1024, 1024, bias: true)
 
-    // MARK: - float
+        // MARK: - float
 
-    let sharded1 = try shardLinear(linear, sharding: .allToSharded, group: world)
-    let sharded2 = try shardLinear(linear, sharding: .shardedToAll, group: world)
+        let sharded1 = try shardLinear(linear, sharding: .allToSharded, group: world)
+        let sharded2 = try shardLinear(linear, sharding: .shardedToAll, group: world)
 
-    let y = linear(x)
-    let y1 = sharded1(x)
-    let y2 = sharded2(x[lower ..< upper, axis: 1])
-    try checkedEval(y, y1, y2)
+        let y = linear(x)
+        let y1 = sharded1(x)
+        let y2 = sharded2(x[lower ..< upper, axis: 1])
+        try checkedEval(y, y1, y2)
 
-    XCTAssertTrue(
-        y.allClose(y2, rtol: 1e-4, atol: 1e-6).item(Bool.self),
-        "sharded-to-all must reproduce the unsharded output")
-    XCTAssertTrue(
-        y[lower ..< upper, axis: 1].allClose(y1, rtol: 1e-4, atol: 1e-6).item(Bool.self),
-        "all-to-sharded must reproduce this rank's slice of the output")
+        XCTAssertTrue(
+            y.allClose(y2, rtol: 1e-4, atol: 1e-6).item(Bool.self),
+            "sharded-to-all must reproduce the unsharded output")
+        XCTAssertTrue(
+            y[lower ..< upper, axis: 1].allClose(y1, rtol: 1e-4, atol: 1e-6).item(Bool.self),
+            "all-to-sharded must reproduce this rank's slice of the output")
 
-    // MARK: - quantized
+        // MARK: - quantized
 
-    // Python: qlin = lin.to_quantized()
-    let quantized = QuantizedLinear(linear)
-    let qSharded1 = try shardLinear(quantized, sharding: .allToSharded, group: world)
-    let qSharded2 = try shardLinear(quantized, sharding: .shardedToAll, group: world)
+        // Python: qlin = lin.to_quantized()
+        let quantized = QuantizedLinear(linear)
+        let qSharded1 = try shardLinear(quantized, sharding: .allToSharded, group: world)
+        let qSharded2 = try shardLinear(quantized, sharding: .shardedToAll, group: world)
 
-    XCTAssertTrue(
-        qSharded1 is QuantizedAllToShardedLinear,
-        "a QuantizedLinear must shard into the quantized layer, not the float one")
-    XCTAssertTrue(qSharded2 is QuantizedShardedToAllLinear)
+        XCTAssertTrue(
+            qSharded1 is QuantizedAllToShardedLinear,
+            "a QuantizedLinear must shard into the quantized layer, not the float one")
+        XCTAssertTrue(qSharded2 is QuantizedShardedToAllLinear)
 
-    let qy = quantized(x)
-    let qy1 = qSharded1(x)
-    let qy2 = qSharded2(x[lower ..< upper, axis: 1])
-    try checkedEval(qy, qy1, qy2)
+        let qy = quantized(x)
+        let qy1 = qSharded1(x)
+        let qy2 = qSharded2(x[lower ..< upper, axis: 1])
+        try checkedEval(qy, qy1, qy2)
 
-    XCTAssertTrue(
-        qy.allClose(qy2, rtol: 1e-4, atol: 1e-6).item(Bool.self),
-        "quantized sharded-to-all must reproduce the unsharded output")
-    XCTAssertTrue(
-        qy[lower ..< upper, axis: 1].allClose(qy1, rtol: 1e-5, atol: 1e-8).item(Bool.self),
-        "quantized all-to-sharded must reproduce this rank's slice")
+        XCTAssertTrue(
+            qy.allClose(qy2, rtol: 1e-4, atol: 1e-6).item(Bool.self),
+            "quantized sharded-to-all must reproduce the unsharded output")
+        XCTAssertTrue(
+            qy[lower ..< upper, axis: 1].allClose(qy1, rtol: 1e-5, atol: 1e-8).item(Bool.self),
+            "quantized all-to-sharded must reproduce this rank's slice")
 
-    // MARK: - a non affine mode, which carries no quantization biases
+        // MARK: - a non affine mode, which carries no quantization biases
 
-    // Python: lin.to_quantized(group_size=32, bits=8, mode="mxfp8")
-    let mxfp8 = QuantizedLinear(linear, groupSize: 32, bits: 8, mode: .mxfp8)
-    XCTAssertEqual(mxfp8.mode, .mxfp8)
-    XCTAssertNil(mxfp8.biases, "mxfp8 quantization has no biases")
+        // Python: lin.to_quantized(group_size=32, bits=8, mode="mxfp8")
+        let mxfp8 = QuantizedLinear(linear, groupSize: 32, bits: 8, mode: .mxfp8)
+        XCTAssertEqual(mxfp8.mode, .mxfp8)
+        XCTAssertNil(mxfp8.biases, "mxfp8 quantization has no biases")
 
-    let mxSharded1 = try XCTUnwrap(
-        try shardLinear(mxfp8, sharding: .allToSharded, group: world) as? QuantizedLinear)
-    let mxSharded2 = try XCTUnwrap(
-        try shardLinear(mxfp8, sharding: .shardedToAll, group: world) as? QuantizedLinear)
+        let mxSharded1 = try XCTUnwrap(
+            try shardLinear(mxfp8, sharding: .allToSharded, group: world) as? QuantizedLinear)
+        let mxSharded2 = try XCTUnwrap(
+            try shardLinear(mxfp8, sharding: .shardedToAll, group: world) as? QuantizedLinear)
 
-    XCTAssertEqual(mxSharded1.mode, .mxfp8, "the mode must survive sharding")
-    XCTAssertEqual(mxSharded2.mode, .mxfp8, "the mode must survive sharding")
-    XCTAssertNil(mxSharded1.biases)
-    XCTAssertNil(mxSharded2.biases)
+        XCTAssertEqual(mxSharded1.mode, .mxfp8, "the mode must survive sharding")
+        XCTAssertEqual(mxSharded2.mode, .mxfp8, "the mode must survive sharding")
+        XCTAssertNil(mxSharded1.biases)
+        XCTAssertNil(mxSharded2.biases)
 
-    let my = mxfp8(x)
-    let my1 = mxSharded1(x)
-    let my2 = mxSharded2(x[lower ..< upper, axis: 1])
-    try checkedEval(my, my1, my2)
+        let my = mxfp8(x)
+        let my1 = mxSharded1(x)
+        let my2 = mxSharded2(x[lower ..< upper, axis: 1])
+        try checkedEval(my, my1, my2)
 
-    XCTAssertTrue(
-        my.allClose(my2, rtol: 1e-4, atol: 1e-6).item(Bool.self),
-        "mxfp8 sharded-to-all must reproduce the unsharded output")
-    XCTAssertTrue(
-        my[lower ..< upper, axis: 1].allClose(my1, rtol: 1e-5, atol: 1e-8).item(Bool.self),
-        "mxfp8 all-to-sharded must reproduce this rank's slice")
+        XCTAssertTrue(
+            my.allClose(my2, rtol: 1e-4, atol: 1e-6).item(Bool.self),
+            "mxfp8 sharded-to-all must reproduce the unsharded output")
+        XCTAssertTrue(
+            my[lower ..< upper, axis: 1].allClose(my1, rtol: 1e-5, atol: 1e-8).item(Bool.self),
+            "mxfp8 all-to-sharded must reproduce this rank's slice")
 
-    // MARK: - quantizing a sharded layer keeps it sharded
+        // MARK: - quantizing a sharded layer keeps it sharded
 
-    let quantizedShard = sharded1.toQuantized(groupSize: 64, bits: 4, mode: .affine)
-    XCTAssertTrue(
-        quantizedShard is QuantizedAllToShardedLinear,
-        "quantizing a sharded layer must not drop the communication")
+        let quantizedShard = sharded1.toQuantized(groupSize: 64, bits: 4, mode: .affine)
+        XCTAssertTrue(
+            quantizedShard is QuantizedAllToShardedLinear,
+            "quantizing a sharded layer must not drop the communication")
 
-    // MARK: - backward
+        // MARK: - backward
 
-    try shardLinearBackward(world: world)
+        try shardLinearBackward(world: world)
+    }
 }
 
 /// The gradients of a sharded model must match the slice of the unsharded
@@ -113,46 +114,47 @@ func shardLinearBody(world: MLXDistributed.Group) throws {
 private func shardLinearBackward(world: MLXDistributed.Group) throws {
     let rank = try world.rank
     let size = try world.size
-    MLXRandom.seed(0xF0F0_F0F0)
-
-    let layers = (0 ..< 4).map { _ in Linear(128, 128) }
-    let model = Sequential(layers: layers[0], layers[1], layers[2], layers[3])
-    let sharded = Sequential(
-        layers: try shardLinear(layers[0], sharding: .allToSharded, group: world),
-        try shardLinear(layers[1], sharding: .shardedToAll, group: world),
-        try shardLinear(layers[2], sharding: .allToSharded, group: world),
-        try shardLinear(layers[3], sharding: .shardedToAll, group: world))
-
-    func loss(_ model: Sequential, _ x: MLXArray, _ y: MLXArray) -> MLXArray {
-        (model(x) * y).sum()
-    }
+    try withRandomState(MLXRandom.RandomState(seed: 0xF0F0_F0F0)) {
+
+        let layers = (0 ..< 4).map { _ in Linear(128, 128) }
+        let model = Sequential(layers: layers[0], layers[1], layers[2], layers[3])
+        let sharded = Sequential(
+            layers: try shardLinear(layers[0], sharding: .allToSharded, group: world),
+            try shardLinear(layers[1], sharding: .shardedToAll, group: world),
+            try shardLinear(layers[2], sharding: .allToSharded, group: world),
+            try shardLinear(layers[3], sharding: .shardedToAll, group: world))
+
+        func loss(_ model: Sequential, _ x: MLXArray, _ y: MLXArray) -> MLXArray {
+            (model(x) * y).sum()
+        }
 
-    let x = MLXRandom.normal([4, 128])
-    let y = MLXRandom.normal([4, 128])
+        let x = MLXRandom.normal([4, 128])
+        let y = MLXRandom.normal([4, 128])
 
-    let (l1, g1) = valueAndGrad(model: model, loss)(model, x, y)
-    let (l2, g2) = valueAndGrad(model: sharded, loss)(sharded, x, y)
-    try checkedEval(l1, g1, l2, g2)
+        let (l1, g1) = valueAndGrad(model: model, loss)(model, x, y)
+        let (l2, g2) = valueAndGrad(model: sharded, loss)(sharded, x, y)
+        try checkedEval(l1, g1, l2, g2)
 
-    let lower = rank * 128 / size
-    let upper = (rank + 1) * 128 / size
-    let full = Dictionary(uniqueKeysWithValues: g1.flattened())
-    let shard = Dictionary(uniqueKeysWithValues: g2.flattened())
+        let lower = rank * 128 / size
+        let upper = (rank + 1) * 128 / size
+        let full = Dictionary(uniqueKeysWithValues: g1.flattened())
+        let shard = Dictionary(uniqueKeysWithValues: g2.flattened())
 
-    XCTAssertTrue(l1.allClose(l2, rtol: 1e-4, atol: 1e-6).item(Bool.self), "loss mismatch")
+        XCTAssertTrue(l1.allClose(l2, rtol: 1e-4, atol: 1e-6).item(Bool.self), "loss mismatch")
 
-    // all-to-sharded splits the output dimension, so rows of the weight
-    // gradient and the bias gradient
-    for layer in [0, 2] {
-        assertGradientShard(full, shard, "layers.\(layer).weight", lower ..< upper, axis: 0)
-        assertGradientShard(full, shard, "layers.\(layer).bias", lower ..< upper, axis: 0)
-    }
+        // all-to-sharded splits the output dimension, so rows of the weight
+        // gradient and the bias gradient
+        for layer in [0, 2] {
+            assertGradientShard(full, shard, "layers.\(layer).weight", lower ..< upper, axis: 0)
+            assertGradientShard(full, shard, "layers.\(layer).bias", lower ..< upper, axis: 0)
+        }
 
-    // sharded-to-all splits the input dimension -- columns of the weight
-    // gradient -- and leaves the bias whole
-    for layer in [1, 3] {
-        assertGradientShard(full, shard, "layers.\(layer).weight", lower ..< upper, axis: 1)
-        assertGradientWhole(full, shard, "layers.\(layer).bias")
+        // sharded-to-all splits the input dimension -- columns of the weight
+        // gradient -- and leaves the bias whole
+        for layer in [1, 3] {
+            assertGradientShard(full, shard, "layers.\(layer).weight", lower ..< upper, axis: 1)
+            assertGradientWhole(full, shard, "layers.\(layer).bias")
+        }
     }
 }
 
@@ -162,40 +164,41 @@ private func shardLinearBackward(world: MLXDistributed.Group) throws {
 /// predicate.  The module aggregates where the sharding requires it, which the
 /// predicate cannot do on its own.
 func shardPredicateBody(world: MLXDistributed.Group) throws {
-    MLXRandom.seed(0xF0F0_F0F0)
-
-    // even layers shard their output channels; odd layers shard their input
-    // channels and keep the bias, which applies to the summed result
-    let sharding: ShardingPredicate = { path, _ in
-        let parts = path.split(separator: ".")
-        guard parts.count > 1, let layer = Int(parts[1]) else { return nil }
-
-        if layer % 2 == 0 {
-            return (0, .count(1))
+    try withRandomState(MLXRandom.RandomState(seed: 0xF0F0_F0F0)) {
+
+        // even layers shard their output channels; odd layers shard their input
+        // channels and keep the bias, which applies to the summed result
+        let sharding: ShardingPredicate = { path, _ in
+            let parts = path.split(separator: ".")
+            guard parts.count > 1, let layer = Int(parts[1]) else { return nil }
+
+            if layer % 2 == 0 {
+                return (0, .count(1))
+            }
+            return parts.last == "bias" ? nil : (-1, .count(1))
         }
-        return parts.last == "bias" ? nil : (-1, .count(1))
-    }
 
-    let model = Sequential(
-        layers: AggregatingConv(3, 128), AggregatingConv(128, 128),
-        AggregatingConv(128, 128), AggregatingConv(128, 3))
-    let sharded = Sequential(
-        layers: AggregatingConv(3, 128),
-        AggregatingConv(128, 128, aggregate: world),
-        AggregatingConv(128, 128),
-        AggregatingConv(128, 3, aggregate: world))
+        let model = Sequential(
+            layers: AggregatingConv(3, 128), AggregatingConv(128, 128),
+            AggregatingConv(128, 128), AggregatingConv(128, 3))
+        let sharded = Sequential(
+            layers: AggregatingConv(3, 128),
+            AggregatingConv(128, 128, aggregate: world),
+            AggregatingConv(128, 128),
+            AggregatingConv(128, 3, aggregate: world))
 
-    _ = sharded.update(parameters: model.parameters())
-    try shardInPlace(sharded, predicate: sharding, group: world)
+        _ = sharded.update(parameters: model.parameters())
+        try shardInPlace(sharded, predicate: sharding, group: world)
 
-    let x = MLXRandom.normal([4, 16, 16, 3])
-    let y1 = model(x)
-    let y2 = sharded(x)
-    try checkedEval(y1, y2)
+        let x = MLXRandom.normal([4, 16, 16, 3])
+        let y1 = model(x)
+        let y2 = sharded(x)
+        try checkedEval(y1, y2)
 
-    XCTAssertTrue(
-        y1.allClose(y2, rtol: 1e-4, atol: 1e-6).item(Bool.self),
-        "a predicate sharded model must reproduce the unsharded output")
+        XCTAssertTrue(
+            y1.allClose(y2, rtol: 1e-4, atol: 1e-6).item(Bool.self),
+            "a predicate sharded model must reproduce the unsharded output")
+    }
 }
 
 /// Sharding swaps the layers of a model in place, the way quantization does.
@@ -205,47 +208,48 @@ func shardPredicateBody(world: MLXDistributed.Group) throws {
 /// because the sharded layers are `Linear` subclasses: a sharded layer that
 /// was not a `Linear` could not be installed at all.
 func moduleSubstitutionBody(world: MLXDistributed.Group) throws {
-    MLXRandom.seed(0xF0F0_F0F0)
-    let x = MLXRandom.normal([4, 64])
-
-    // the check has to be able to fail: a module that is not a Linear is refused
-    let control = ProjectionBlock(64, 256)
-    let notLinear: [(String, Module)] = [("up", Identity()), ("down", control.down)]
-    XCTAssertThrowsError(
-        try control.update(modules: ModuleChildren.unflattened(notLinear), verify: .all)
-    ) { error in
-        XCTAssertTrue(error is UpdateError, "unexpected error: \(error)")
-    }
+    try withRandomState(MLXRandom.RandomState(seed: 0xF0F0_F0F0)) {
+        let x = MLXRandom.normal([4, 64])
 
-    // a float model, and one quantized first like a loaded quantized model
-    for quantized in [false, true] {
-        let block = ProjectionBlock(64, 256)
-        if quantized {
-            quantize(model: block)
-        }
-        let expected = block(x)
-        try checkedEval(expected)
-
-        // what a model's shard method does
-        let up = try shardLinear(block.up, sharding: .allToSharded, group: world)
-        let down = try shardLinear(block.down, sharding: .shardedToAll, group: world)
-        let sharded: [(String, Module)] = [("up", up), ("down", down)]
-        try block.update(modules: ModuleChildren.unflattened(sharded), verify: .all)
-
-        if quantized {
-            XCTAssertTrue(block.up is QuantizedAllToShardedLinear, "\(type(of: block.up))")
-            XCTAssertTrue(block.down is QuantizedShardedToAllLinear, "\(type(of: block.down))")
-        } else {
-            XCTAssertTrue(block.up is AllToShardedLinear, "\(type(of: block.up))")
-            XCTAssertTrue(block.down is ShardedToAllLinear, "\(type(of: block.down))")
+        // the check has to be able to fail: a module that is not a Linear is refused
+        let control = ProjectionBlock(64, 256)
+        let notLinear: [(String, Module)] = [("up", Identity()), ("down", control.down)]
+        XCTAssertThrowsError(
+            try control.update(modules: ModuleChildren.unflattened(notLinear), verify: .all)
+        ) { error in
+            XCTAssertTrue(error is UpdateError, "unexpected error: \(error)")
         }
 
-        // every process computes the output of the whole block
-        let y = block(x)
-        try checkedEval(y)
-        XCTAssertTrue(
-            expected.allClose(y, rtol: 1e-4, atol: 1e-6).item(Bool.self),
-            "the sharded block must reproduce the unsharded output (quantized: \(quantized))")
+        // a float model, and one quantized first like a loaded quantized model
+        for quantized in [false, true] {
+            let block = ProjectionBlock(64, 256)
+            if quantized {
+                quantize(model: block)
+            }
+            let expected = block(x)
+            try checkedEval(expected)
+
+            // what a model's shard method does
+            let up = try shardLinear(block.up, sharding: .allToSharded, group: world)
+            let down = try shardLinear(block.down, sharding: .shardedToAll, group: world)
+            let sharded: [(String, Module)] = [("up", up), ("down", down)]
+            try block.update(modules: ModuleChildren.unflattened(sharded), verify: .all)
+
+            if quantized {
+                XCTAssertTrue(block.up is QuantizedAllToShardedLinear, "\(type(of: block.up))")
+                XCTAssertTrue(block.down is QuantizedShardedToAllLinear, "\(type(of: block.down))")
+            } else {
+                XCTAssertTrue(block.up is AllToShardedLinear, "\(type(of: block.up))")
+                XCTAssertTrue(block.down is ShardedToAllLinear, "\(type(of: block.down))")
+            }
+
+            // every process computes the output of the whole block
+            let y = block(x)
+            try checkedEval(y)
+            XCTAssertTrue(
+                expected.allClose(y, rtol: 1e-4, atol: 1e-6).item(Bool.self),
+                "the sharded block must reproduce the unsharded output (quantized: \(quantized))")
+        }
     }
 }
 
@@ -329,54 +333,55 @@ func shardingEdgeCasesBody(world: MLXDistributed.Group) throws {
     // process has to hold the same bias.  Seed the processes differently:
     // nothing keeps their generators in step -- they may have done different
     // work before this -- and the bias must not depend on that.
-    MLXRandom.seed(UInt64(rank + 1))
-    let layer = try ShardedToAllLinear(8 * size, 4, group: world)
-    let bias = try XCTUnwrap(layer.bias)
-    let gathered = MLXDistributed.allGather(bias, group: world)
-    try checkedEval(gathered)
-
-    let width = bias.dim(0)
-    for rank in 1 ..< size {
-        XCTAssertTrue(
-            gathered[0 ..< width].allClose(gathered[rank * width ..< (rank + 1) * width])
-                .item(Bool.self),
-            "rank \(rank) holds a different bias than rank 0")
-    }
+    try withRandomState(MLXRandom.RandomState(seed: UInt64(rank + 1))) {
+        let layer = try ShardedToAllLinear(8 * size, 4, group: world)
+        let bias = try XCTUnwrap(layer.bias)
+        let gathered = MLXDistributed.allGather(bias, group: world)
+        try checkedEval(gathered)
+
+        let width = bias.dim(0)
+        for rank in 1 ..< size {
+            XCTAssertTrue(
+                gathered[0 ..< width].allClose(gathered[rank * width ..< (rank + 1) * width])
+                    .item(Bool.self),
+                "rank \(rank) holds a different bias than rank 0")
+        }
 
-    // A dimension can divide by the group size while one of its segments does
-    // not.  That has to be reported rather than trap inside split(parts:).
-    if size > 1 {
+        // A dimension can divide by the group size while one of its segments does
+        // not.  That has to be reported rather than trap inside split(parts:).
+        if size > 1 {
+            XCTAssertThrowsError(
+                try shardLinear(
+                    Linear(8 * size, 8), sharding: .shardedToAll,
+                    segments: .indices([3]), group: world)
+            ) { error in
+                XCTAssertEqual(
+                    error as? ShardingError,
+                    .indivisible(dimension: "segment of weight", of: 3, across: size))
+            }
+        }
+
+        // Segments that cannot split a dimension at all are reported too, in a
+        // group of one as well.  split(parts:) would call fatalError.
         XCTAssertThrowsError(
             try shardLinear(
-                Linear(8 * size, 8), sharding: .shardedToAll,
-                segments: .indices([3]), group: world)
+                Linear(8 * size, 8), sharding: .shardedToAll, segments: .count(3),
+                group: world)
         ) { error in
             XCTAssertEqual(
-                error as? ShardingError,
-                .indivisible(dimension: "segment of weight", of: 3, across: size))
+                error as? ShardingError, .invalidSegments(.count(3), dimension: 8 * size))
         }
-    }
 
-    // Segments that cannot split a dimension at all are reported too, in a
-    // group of one as well.  split(parts:) would call fatalError.
-    XCTAssertThrowsError(
-        try shardLinear(
-            Linear(8 * size, 8), sharding: .shardedToAll, segments: .count(3),
-            group: world)
-    ) { error in
-        XCTAssertEqual(
-            error as? ShardingError, .invalidSegments(.count(3), dimension: 8 * size))
-    }
-
-    // A QuantizedLinear is a Linear, so the float layers accept one as far as
-    // the compiler is concerned, but they would keep its packed weight as a
-    // float weight and drop its scales.
-    let quantized = QuantizedLinear(Linear(64, 64))
-    XCTAssertThrowsError(try AllToShardedLinear(quantized, group: world)) { error in
-        XCTAssertEqual(error as? ShardingError, .quantizedLayer)
-    }
-    XCTAssertThrowsError(try ShardedToAllLinear(quantized, group: world)) { error in
-        XCTAssertEqual(error as? ShardingError, .quantizedLayer)
+        // A QuantizedLinear is a Linear, so the float layers accept one as far as
+        // the compiler is concerned, but they would keep its packed weight as a
+        // float weight and drop its scales.
+        let quantized = QuantizedLinear(Linear(64, 64))
+        XCTAssertThrowsError(try AllToShardedLinear(quantized, group: world)) { error in
+            XCTAssertEqual(error as? ShardingError, .quantizedLayer)
+        }
+        XCTAssertThrowsError(try ShardedToAllLinear(quantized, group: world)) { error in
+            XCTAssertEqual(error as? ShardingError, .quantizedLayer)
+        }
     }
 
     // A quantized layer packs its inputs in the weight and groups them in the
@@ -384,35 +389,37 @@ func shardingEdgeCasesBody(world: MLXDistributed.Group) throws {
     // each.  Scaled through a Double, 3840 of 5632 inputs came to
     // 479.99999999999994 packed words and 59.99999999999999 groups, which
     // truncate to the wrong positions.  The shard has to reproduce the layer.
-    MLXRandom.seed(0xF0F0_F0F0)
-    let inputs = 5632
-    let boundary = 3840
-    let wide = QuantizedLinear(Linear(inputs, 8))
-    let segmented = try QuantizedShardedToAllLinear(
-        wide, segments: .indices([boundary]), group: world)
-
-    // this rank's part of each segment
-    let first = (rank * boundary / size) ..< ((rank + 1) * boundary / size)
-    let second =
-        (boundary + rank * (inputs - boundary) / size)
-        ..< (boundary + (rank + 1) * (inputs - boundary) / size)
-
-    let x = MLXRandom.normal([2, inputs])
-    let y = wide(x)
-    let ySharded = segmented(concatenated([x[first, axis: 1], x[second, axis: 1]], axis: 1))
-    try checkedEval(y, ySharded)
+    try withRandomState(MLXRandom.RandomState(seed: 0xF0F0_F0F0)) {
+        let inputs = 5632
+        let boundary = 3840
+        let wide = QuantizedLinear(Linear(inputs, 8))
+        let segmented = try QuantizedShardedToAllLinear(
+            wide, segments: .indices([boundary]), group: world)
+
+        // this rank's part of each segment
+        let first = (rank * boundary / size) ..< ((rank + 1) * boundary / size)
+        let second =
+            (boundary + rank * (inputs - boundary) / size)
+            ..< (boundary + (rank + 1) * (inputs - boundary) / size)
+
+        let x = MLXRandom.normal([2, inputs])
+        let y = wide(x)
+        let ySharded = segmented(
+            concatenated([x[first, axis: 1], x[second, axis: 1]], axis: 1))
+        try checkedEval(y, ySharded)
 
-    XCTAssertTrue(
-        y.allClose(ySharded, rtol: 1e-4, atol: 1e-6).item(Bool.self),
-        "a segmented quantized shard must reproduce the unsharded output")
+        XCTAssertTrue(
+            y.allClose(ySharded, rtol: 1e-4, atol: 1e-6).item(Bool.self),
+            "a segmented quantized shard must reproduce the unsharded output")
 
-    // a boundary inside a quantization group cannot be sharded at all
-    XCTAssertThrowsError(
-        try QuantizedShardedToAllLinear(wide, segments: .indices([272]), group: world)
-    ) { error in
-        XCTAssertEqual(
-            error as? ShardingError,
-            .quantizationGroup(inputDimensions: 272, groupSize: 64, across: size))
+        // a boundary inside a quantization group cannot be sharded at all
+        XCTAssertThrowsError(
+            try QuantizedShardedToAllLinear(wide, segments: .indices([272]), group: world)
+        ) { error in
+            XCTAssertEqual(
+                error as? ShardingError,
+                .quantizationGroup(inputDimensions: 272, groupSize: 64, across: size))
+        }
     }
 }
 
@@ -599,29 +606,30 @@ private func assertGradientWhole(
 func quantizedShardInPlaceBody(world: MLXDistributed.Group) throws {
     let rank = try world.rank
     let size = try world.size
-    MLXRandom.seed(0)
-    let x = MLXRandom.normal([4, 1024])
-    let model = Sequential(layers: QuantizedLinear(Linear(1024, 32, bias: false)))
-    let expected = model(x)
-
-    try shardInPlace(model, sharding: .shardedToAll, segments: .indices([512]), group: world)
-
-    // each process holds its slice of both segments; the sum of the partial
-    // results is the unsharded result
-    let step = 512 / size
-    let lower = rank * step
-    let xPart = concatenated(
-        [x[0..., lower ..< lower + step], x[0..., 512 + lower ..< 512 + lower + step]], axis: 1)
-    let result = MLXDistributed.allSum(model(xPart), group: world)
-    XCTAssertTrue(
-        expected.allClose(result, rtol: 1e-4, atol: 1e-6).item(Bool.self),
-        "integer segments landed in the wrong place")
-
-    // a segment has to hold whole quantization groups on every process
-    XCTAssertThrowsError(
-        try shardInPlace(
-            QuantizedLinear(Linear(1024, 32)), sharding: .shardedToAll,
-            segments: .indices([528]), group: world))
+    try withRandomState(MLXRandom.RandomState(seed: 0)) {
+        let x = MLXRandom.normal([4, 1024])
+        let model = Sequential(layers: QuantizedLinear(Linear(1024, 32, bias: false)))
+        let expected = model(x)
+
+        try shardInPlace(model, sharding: .shardedToAll, segments: .indices([512]), group: world)
+
+        // each process holds its slice of both segments; the sum of the partial
+        // results is the unsharded result
+        let step = 512 / size
+        let lower = rank * step
+        let xPart = concatenated(
+            [x[0..., lower ..< lower + step], x[0..., 512 + lower ..< 512 + lower + step]], axis: 1)
+        let result = MLXDistributed.allSum(model(xPart), group: world)
+        XCTAssertTrue(
+            expected.allClose(result, rtol: 1e-4, atol: 1e-6).item(Bool.self),
+            "integer segments landed in the wrong place")
+
+        // a segment has to hold whole quantization groups on every process
+        XCTAssertThrowsError(
+            try shardInPlace(
+                QuantizedLinear(Linear(1024, 32)), sharding: .shardedToAll,
+                segments: .indices([528]), group: world))
+    }
 }
 
 /// The single process half: sharding is the identity in a group of size one,
diff --git a/Tests/MLXTests/ErrorTests.swift b/Tests/MLXTests/ErrorTests.swift
index 320407d5..41abacaf 100644
--- a/Tests/MLXTests/ErrorTests.swift
+++ b/Tests/MLXTests/ErrorTests.swift
@@ -2,7 +2,7 @@
 
 import Foundation
 import MLX
-import XCTest
+import Testing
 
 private func triggerError() {
     let a = MLXArray(0 ..< 10, [2, 5])
@@ -14,9 +14,10 @@ private func triggerError() {
     eval(a + b)
 }
 
-class ErrorTests: XCTestCase {
+@Suite
+struct ErrorTests {
 
-    func testErrorHandler() {
+    @Test func errorHandler() throws {
         // a test with a custom error handler that captures the first error seen
         final class BoxedError: @unchecked (Sendable) {
             let lock = NSLock()
@@ -46,68 +47,49 @@ class ErrorTests: XCTestCase {
             triggerError()
         }
 
-        if let error = boxedError.error {
-            let prefix = error.split(separator: ". at")[0]
-            XCTAssertEqual(prefix, "[broadcast_shapes] Shapes (2,5) and (3,5) cannot be broadcast")
-        } else {
-            XCTFail("boxedError.error should be set")
-        }
+        let error = try #require(boxedError.error, "boxedError.error should be set")
+        let prefix = error.split(separator: ". at")[0]
+        #expect(prefix == "[broadcast_shapes] Shapes (2,5) and (3,5) cannot be broadcast")
     }
 
-    func testWithErrorCheck() {
-        do {
+    @Test func withErrorCheck() {
+        #expect(throws: MLXError.self) {
             try withError { error in
                 triggerError()
                 try error.check()
             }
-            XCTFail("should throw")
-        } catch {
-            print(error)
         }
     }
 
-    func testWithErrorThrow() {
-        do {
-            // outer container should throw
+    @Test func withErrorThrow() {
+        // outer container should throw
+        #expect(throws: MLXError.self) {
             try withError {
                 triggerError()
             }
-            XCTFail("should throw")
-        } catch {
-            print(error)
         }
     }
 
-    func testWithErrorThrowAsync() async {
-        do {
+    @Test func withErrorThrowAsync() async {
+        await #expect(throws: MLXError.self) {
             try await withError {
                 let t = Task {
                     triggerError()
                 }
                 await t.value
             }
-            XCTFail("should throw")
-        } catch {
-            print(error)
         }
     }
 
-    func testWithErrorThrowNested() {
-        do {
-            // outer container should throw
-            try withError {
-                do {
-                    try withError {
-                        triggerError()
-                    }
-                    XCTFail("should throw")
-                } catch {
-                    // expected
+    @Test func withErrorThrowNested() throws {
+        // the outer layer does not produce an error -- if it throws the test fails
+        let _ = try withError {
+            // the inner container should throw
+            #expect(throws: MLXError.self) {
+                try withError {
+                    triggerError()
                 }
-                // the outer layer does not produce an error
             }
-        } catch {
-            XCTFail("unexpected outer error: \(error)")
         }
     }
 
diff --git a/Tests/MLXTests/ExportTests.swift b/Tests/MLXTests/ExportTests.swift
index babcb56b..632afd8c 100644
--- a/Tests/MLXTests/ExportTests.swift
+++ b/Tests/MLXTests/ExportTests.swift
@@ -2,27 +2,32 @@
 
 import Foundation
 import MLX
-import XCTest
+import Testing
 
-class ExportTests: XCTestCase {
+/// A class rather than a struct so that `deinit` can clean up -- Swift Testing
+/// creates a new instance per test, so `init` / `deinit` take the place of
+/// XCTest's `setUpWithError` / `tearDownWithError`.
+@Suite
+final class ExportTests {
 
     let temporaryPath = FileManager.default.temporaryDirectory.appending(
         path: UUID().uuidString,
         directoryHint: .isDirectory
     )
 
-    override func setUpWithError() throws {
+    init() throws {
         try FileManager.default.createDirectory(
             at: temporaryPath,
             withIntermediateDirectories: false
         )
     }
 
-    override func tearDownWithError() throws {
-        try FileManager.default.removeItem(at: temporaryPath)
+    deinit {
+        // deinit cannot throw -- a leftover temporary directory is not a test failure
+        try? FileManager.default.removeItem(at: temporaryPath)
     }
 
-    func testExportFunction() throws {
+    @Test func exportAndImportFunction() throws {
         let url = temporaryPath.appending(path: "fn.mlxfn")
 
         func f(arrays: [MLXArray]) -> [MLXArray] {
@@ -42,10 +47,10 @@ class ExportTests: XCTestCase {
 
         // call it -- the shapes and labels have to match
         let r = try f2(a, y: b)[0]
-        assertEqual(r, MLXArray([50, 100, 200]))
+        expectEqual(r, MLXArray([50, 100, 200]))
     }
 
-    func testExportFunctions() throws {
+    @Test func exportAndImportShapelessFunctions() throws {
         let url = temporaryPath.appending(path: "fn.mlxfn")
 
         func f(_ arrays: [MLXArray]) -> [MLXArray] {
@@ -68,16 +73,16 @@ class ExportTests: XCTestCase {
         let c = MLXArray([1, 2, 3])
 
         let r1 = try f2(a)[0]
-        assertEqual(r1, a)
+        expectEqual(r1, a)
 
         let r2 = try f2(a, b)[0]
-        assertEqual(r2, a + b)
+        expectEqual(r2, a + b)
 
         let r3 = try f2(a, b, c)[0]
-        assertEqual(r3, a + b + c)
+        expectEqual(r3, a + b + c)
     }
 
-    func testExportError() {
+    @Test func exportToInvalidPathThrows() {
         func f(arrays: [MLXArray]) -> [MLXArray] {
             [arrays[0] * arrays[1]]
         }
@@ -85,11 +90,8 @@ class ExportTests: XCTestCase {
         let x = MLXArray(1)
         let y = MLXArray([1, 2, 3])
 
-        do {
+        #expect(throws: MLXError.self) {
             try exportFunction(to: URL(fileURLWithPath: "/does/not/exist"), f)(x, y: y)
-            XCTFail("should throw")
-        } catch {
-            // expected
         }
     }
 
diff --git a/Tests/MLXTests/FFTTests.swift b/Tests/MLXTests/FFTTests.swift
index 258b59c5..9423ec26 100644
--- a/Tests/MLXTests/FFTTests.swift
+++ b/Tests/MLXTests/FFTTests.swift
@@ -2,23 +2,21 @@
 
 import Foundation
 import MLX
-import XCTest
+import Testing
 
-class FFTTests: XCTestCase {
+@Suite
+struct FFTTests {
 
-    override class func setUp() {
-    }
-
-    func testFFTRoundTrip() {
+    @Test func fftRoundTrip() {
         let a = MLXArray(converting: [1.0, 2.0, 3.0, 4.0])
 
         for norm in [FFTNorm.backward, .ortho, .forward] {
             let r = ifft(fft(a, norm: norm), norm: norm)
-            assertEqual(r.realPart(), a, atol: 1e-5)
+            expectEqual(r.realPart(), a, atol: 1e-5)
         }
     }
 
-    func testFFTNormScaling() {
+    @Test func fftNormScaling() {
         let a = MLXArray(converting: [1.0, 1.0, 1.0, 1.0])
         let n = 4.0
 
@@ -28,45 +26,45 @@ class FFTTests: XCTestCase {
         let ortho = fft(a, norm: .ortho)[0].realPart().item(Double.self)
         let forward = fft(a, norm: .forward)[0].realPart().item(Double.self)
 
-        XCTAssertEqual(backward, n, accuracy: 1e-5)
-        XCTAssertEqual(ortho, n / n.squareRoot(), accuracy: 1e-5)
-        XCTAssertEqual(forward, 1.0, accuracy: 1e-5)
+        expectEqual(backward, n, accuracy: 1e-5)
+        expectEqual(ortho, n / n.squareRoot(), accuracy: 1e-5)
+        expectEqual(forward, 1.0, accuracy: 1e-5)
     }
 
-    func testFFTNormDefaultIsBackward() {
+    @Test func fftNormDefaultIsBackward() {
         let a = MLXArray(converting: [1.0, 2.0, 3.0, 4.0])
-        assertEqual(fft(a).realPart(), fft(a, norm: .backward).realPart(), atol: 1e-5)
+        expectEqual(fft(a).realPart(), fft(a, norm: .backward).realPart(), atol: 1e-5)
     }
 
-    func testFFTNAndAxes() {
+    @Test func fftNAndAxes() {
         let a = MLXArray(0 ..< 24, [2, 3, 4]).asType(.float32)
 
         // norm should thread through each of the n-dimensional variants
         for norm in [FFTNorm.backward, .ortho, .forward] {
-            assertEqual(ifftn(fftn(a, norm: norm), norm: norm).realPart(), a, atol: 1e-4)
-            assertEqual(
+            expectEqual(ifftn(fftn(a, norm: norm), norm: norm).realPart(), a, atol: 1e-4)
+            expectEqual(
                 ifft2(fft2(a, norm: norm), norm: norm).realPart(), a, atol: 1e-4)
-            assertEqual(
+            expectEqual(
                 irfftn(rfftn(a, norm: norm), norm: norm), a, atol: 1e-4)
         }
     }
 
-    func testFFTFrequencies() {
+    @Test func fftFrequencies() {
         // matches numpy/mlx: [0, 1, 2, 3, -4, -3, -2, -1] / (n * d)
         let f = fftfreq(8)
-        assertEqual(
+        expectEqual(
             f, MLXArray(converting: [0, 0.125, 0.25, 0.375, -0.5, -0.375, -0.25, -0.125]),
             atol: 1e-6)
 
         // d scales the result
         let scaled = fftfreq(8, d: 0.5)
-        assertEqual(scaled, f * 2, atol: 1e-6)
+        expectEqual(scaled, f * 2, atol: 1e-6)
     }
 
-    func testRFFTFrequencies() {
+    @Test func rfftFrequencies() {
         // only the non-negative frequencies -- n / 2 + 1 of them
         let f = rfftfreq(8)
-        XCTAssertEqual(f.shape, [5])
-        assertEqual(f, MLXArray(converting: [0, 0.125, 0.25, 0.375, 0.5]), atol: 1e-6)
+        #expect(f.shape == [5])
+        expectEqual(f, MLXArray(converting: [0, 0.125, 0.25, 0.375, 0.5]), atol: 1e-6)
     }
 }
diff --git a/Tests/MLXTests/FP8Tests.swift b/Tests/MLXTests/FP8Tests.swift
index 9df64346..1f29c694 100644
--- a/Tests/MLXTests/FP8Tests.swift
+++ b/Tests/MLXTests/FP8Tests.swift
@@ -2,39 +2,37 @@
 
 import Foundation
 import MLX
-import XCTest
+import Testing
 
-class FP8Tests: XCTestCase {
+@Suite
+struct FP8Tests {
 
-    override class func setUp() {
-    }
-
-    func testFromFP8DecodesE4M3BytesPositive() {
+    @Test func fromFP8DecodesE4M3BytesPositive() {
         // 127 and 255 are NaN and don't compare equal, so skip them
         let allBytes = (0 ... 126).map { UInt8($0) }
         let bytes = MLXArray(allBytes)
         let decoded = fromFP8(bytes, dtype: .float32)
 
-        XCTAssertEqual(decoded.dtype, .float32)
-        XCTAssertEqual(decoded.shape, [127])
+        #expect(decoded.dtype == .float32)
+        #expect(decoded.shape == [127])
 
         let expected = MLXArray(allBytes.map { e4m3Value($0) })
-        assertEqual(decoded, expected)
+        expectEqual(decoded, expected)
     }
 
-    func testFromFP8DecodesE4M3BytesNegative() {
+    @Test func fromFP8DecodesE4M3BytesNegative() {
         let allBytes = (128 ... 254).map { UInt8($0) }
         let bytes = MLXArray(allBytes)
         let decoded = fromFP8(bytes, dtype: .float32)
 
-        XCTAssertEqual(decoded.dtype, .float32)
-        XCTAssertEqual(decoded.shape, [127])
+        #expect(decoded.dtype == .float32)
+        #expect(decoded.shape == [127])
 
         let expected = MLXArray(allBytes.map { e4m3Value($0) })
-        assertEqual(decoded, expected)
+        expectEqual(decoded, expected)
     }
 
-    func testToFP8EncodesE4M3Bytes() {
+    @Test func toFP8EncodesE4M3Bytes() {
         let values = MLXArray(
             [
                 0.0, -0.0, 1.0, -2.0, 0.125, 0.015625, 448.0, -448.0,
@@ -42,11 +40,10 @@ class FP8Tests: XCTestCase {
 
         let encoded = toFP8(values)
 
-        XCTAssertEqual(encoded.dtype, .uint8)
-        XCTAssertEqual(encoded.shape, values.shape)
-        XCTAssertEqual(
-            encoded.asArray(UInt8.self),
-            [0x00, 0x80, 0x38, 0xC0, 0x20, 0x08, 0x7E, 0xFE])
+        #expect(encoded.dtype == .uint8)
+        #expect(encoded.shape == values.shape)
+        #expect(
+            encoded.asArray(UInt8.self) == [0x00, 0x80, 0x38, 0xC0, 0x20, 0x08, 0x7E, 0xFE])
     }
 }
 
diff --git a/Tests/MLXTests/GraphUtilsTests.swift b/Tests/MLXTests/GraphUtilsTests.swift
index 81c7c497..d36fa1cd 100644
--- a/Tests/MLXTests/GraphUtilsTests.swift
+++ b/Tests/MLXTests/GraphUtilsTests.swift
@@ -2,34 +2,32 @@
 
 import Foundation
 import MLX
-import XCTest
+import Testing
 
-class GraphUtilsTests: XCTestCase {
+@Suite
+struct GraphUtilsTests {
 
-    override class func setUp() {
-    }
-
-    func testExportToDot() {
+    @Test func exportToDotIncludesPrimitives() {
         let x = MLXArray([1, 2, 3])
         let dot = exportToDot([x * 2 + 1])
 
-        XCTAssertFalse(dot.isEmpty)
-        XCTAssertTrue(dot.contains("digraph"), "expected a digraph, got:\n\(dot)")
+        #expect(!dot.isEmpty)
+        #expect(dot.contains("digraph"), "expected a digraph, got:\n\(dot)")
 
         // the graph should mention the primitives that produce the output
-        XCTAssertTrue(dot.contains("Add"), "expected Add in:\n\(dot)")
-        XCTAssertTrue(dot.contains("Multiply"), "expected Multiply in:\n\(dot)")
+        #expect(dot.contains("Add"), "expected Add in:\n\(dot)")
+        #expect(dot.contains("Multiply"), "expected Multiply in:\n\(dot)")
     }
 
-    func testGraphDescription() {
+    @Test func graphDescriptionIncludesPrimitives() {
         let x = MLXArray([1, 2, 3])
         let description = graphDescription([x + x])
 
-        XCTAssertFalse(description.isEmpty)
-        XCTAssertTrue(description.contains("Add"), "expected Add in:\n\(description)")
+        #expect(!description.isEmpty)
+        #expect(description.contains("Add"), "expected Add in:\n\(description)")
     }
 
-    func testNodeNamer() {
+    @Test func nodeNamer() {
         let x = MLXArray([1, 2, 3])
         let y = MLXArray([4, 5, 6])
 
@@ -37,25 +35,25 @@ class GraphUtilsTests: XCTestCase {
         namer.setName("lhs", for: x)
         namer.setName("rhs", for: y)
 
-        XCTAssertEqual(namer.name(for: x), "lhs")
-        XCTAssertEqual(namer.name(for: y), "rhs")
+        #expect(namer.name(for: x) == "lhs")
+        #expect(namer.name(for: y) == "rhs")
 
         let dot = exportToDot([x + y], namer: namer)
-        XCTAssertTrue(dot.contains("lhs"), "expected lhs in:\n\(dot)")
-        XCTAssertTrue(dot.contains("rhs"), "expected rhs in:\n\(dot)")
+        #expect(dot.contains("lhs"), "expected lhs in:\n\(dot)")
+        #expect(dot.contains("rhs"), "expected rhs in:\n\(dot)")
     }
 
-    func testNodeNamerGeneratesNames() {
+    @Test func nodeNamerGeneratesNames() {
         // unnamed arrays get a generated name, and it is stable
         let x = MLXArray([1, 2, 3])
         let namer = NodeNamer()
 
         let first = namer.name(for: x)
-        XCTAssertNotNil(first)
-        XCTAssertEqual(namer.name(for: x), first)
+        #expect(first != nil)
+        #expect(namer.name(for: x) == first)
     }
 
-    func testGraphIsNotEvaluated() {
+    @Test func graphIsNotEvaluated() {
         // describing the graph should not force evaluation
         let x = MLXArray([1, 2, 3])
         let y = x * 2
@@ -63,6 +61,6 @@ class GraphUtilsTests: XCTestCase {
         _ = exportToDot([y])
 
         // still usable afterwards
-        assertEqual(y, MLXArray([2, 4, 6]))
+        expectEqual(y, MLXArray([2, 4, 6]))
     }
 }
diff --git a/Tests/MLXTests/LinalgTests.swift b/Tests/MLXTests/LinalgTests.swift
index d0d0856f..92eebfbc 100644
--- a/Tests/MLXTests/LinalgTests.swift
+++ b/Tests/MLXTests/LinalgTests.swift
@@ -2,65 +2,63 @@
 
 import Foundation
 import MLX
-import XCTest
+import Testing
 
-class LinalgTests: XCTestCase {
+@Suite
+struct LinalgTests {
 
-    override class func setUp() {
-    }
-
-    func testNormNoAxes() {
+    @Test func normNoAxes() {
         let a = MLXArray(0 ..< 9) - 4
         let b = a.reshaped(3, 3)
 
-        XCTAssertEqual(MLXLinalg.norm(a).item(Float.self), 7.74597, accuracy: 0.001)
-        XCTAssertEqual(MLXLinalg.norm(b).item(Float.self), 7.74597, accuracy: 0.001)
+        expectEqual(MLXLinalg.norm(a).item(Float.self), 7.74597, accuracy: 0.001)
+        expectEqual(MLXLinalg.norm(b).item(Float.self), 7.74597, accuracy: 0.001)
 
-        XCTAssertEqual(MLXLinalg.norm(b, ord: .fro).item(Float.self), 7.74597, accuracy: 0.001)
+        expectEqual(MLXLinalg.norm(b, ord: .fro).item(Float.self), 7.74597, accuracy: 0.001)
 
         // Double.infinity
-        XCTAssertEqual(MLXLinalg.norm(a, ord: .infinity).item(Float.self), 4, accuracy: 0.001)
-        XCTAssertEqual(MLXLinalg.norm(b, ord: .infinity).item(Float.self), 9, accuracy: 0.001)
+        expectEqual(MLXLinalg.norm(a, ord: .infinity).item(Float.self), 4, accuracy: 0.001)
+        expectEqual(MLXLinalg.norm(b, ord: .infinity).item(Float.self), 9, accuracy: 0.001)
 
-        XCTAssertEqual(MLXLinalg.norm(a, ord: -.infinity).item(Float.self), 0, accuracy: 0.001)
-        XCTAssertEqual(MLXLinalg.norm(b, ord: -.infinity).item(Float.self), 2, accuracy: 0.001)
+        expectEqual(MLXLinalg.norm(a, ord: -.infinity).item(Float.self), 0, accuracy: 0.001)
+        expectEqual(MLXLinalg.norm(b, ord: -.infinity).item(Float.self), 2, accuracy: 0.001)
 
-        XCTAssertEqual(MLXLinalg.norm(a, ord: 1).item(Float.self), 20, accuracy: 0.001)
-        XCTAssertEqual(MLXLinalg.norm(b, ord: 1).item(Float.self), 7, accuracy: 0.001)
+        expectEqual(MLXLinalg.norm(a, ord: 1).item(Float.self), 20, accuracy: 0.001)
+        expectEqual(MLXLinalg.norm(b, ord: 1).item(Float.self), 7, accuracy: 0.001)
 
-        XCTAssertEqual(MLXLinalg.norm(a, ord: -1).item(Float.self), 0, accuracy: 0.001)
-        XCTAssertEqual(MLXLinalg.norm(b, ord: -1).item(Float.self), 6, accuracy: 0.001)
+        expectEqual(MLXLinalg.norm(a, ord: -1).item(Float.self), 0, accuracy: 0.001)
+        expectEqual(MLXLinalg.norm(b, ord: -1).item(Float.self), 6, accuracy: 0.001)
     }
 
-    func testNormAxis() {
+    @Test func normAxis() {
         let c = MLXArray([1, 2, 3, -1, 1, 4], [2, 3])
 
-        assertEqual(MLXLinalg.norm(c, axis: 0), MLXArray(converting: [1.41421, 2.23607, 5]))
-        assertEqual(MLXLinalg.norm(c, ord: 1, axis: 1), MLXArray(converting: [6, 6]))
+        expectEqual(MLXLinalg.norm(c, axis: 0), MLXArray(converting: [1.41421, 2.23607, 5]))
+        expectEqual(MLXLinalg.norm(c, ord: 1, axis: 1), MLXArray(converting: [6, 6]))
     }
 
-    func testNormAxes() {
+    @Test func normAxes() {
         let m = MLXArray(0 ..< 8, [2, 2, 2])
 
-        assertEqual(MLXLinalg.norm(m, axes: [1, 2]), MLXArray(converting: [3.74166, 11.225]))
+        expectEqual(MLXLinalg.norm(m, axes: [1, 2]), MLXArray(converting: [3.74166, 11.225]))
     }
 
-    func testQR() {
+    @Test func qr() {
         let a = MLXArray(converting: [2, 3, 1, 2], [2, 2])
         let (q, r) = MLXLinalg.qr(a, stream: .cpu)
 
-        assertEqual(q, MLXArray(converting: [-0.894427, -0.447214, -0.447214, 0.894427], [2, 2]))
-        assertEqual(r, MLXArray(converting: [-2.23607, -3.57771, 0, 0.447214], [2, 2]))
+        expectEqual(q, MLXArray(converting: [-0.894427, -0.447214, -0.447214, 0.894427], [2, 2]))
+        expectEqual(r, MLXArray(converting: [-2.23607, -3.57771, 0, 0.447214], [2, 2]))
     }
 
-    func testSVDOverload() {
+    @Test func svdOverload() {
         let a = MLXRandom.uniform(0 ..< 1, [10, 10])
 
         Stream.withNewDefaultStream(device: .cpu) {
             let (_, s, _) = MLXLinalg.svd(a)
             let s2 = MLXLinalg.svd(a)
 
-            XCTAssertEqual(s.shape, s2.shape)
+            #expect(s.shape == s2.shape)
         }
     }
 }
diff --git a/Tests/MLXTests/LossTests.swift b/Tests/MLXTests/LossTests.swift
index 56616657..3c364cfc 100644
--- a/Tests/MLXTests/LossTests.swift
+++ b/Tests/MLXTests/LossTests.swift
@@ -3,13 +3,12 @@
 import Foundation
 import MLX
 import MLXNN
-import XCTest
+import Testing
 
-class LossTests: XCTestCase {
-    override class func setUp() {
-    }
+@Suite
+struct LossTests {
 
-    func testCrossEntropy() {
+    @Test func crossEntropyIndicesAndProbabilities() {
         // This is just testing that crossEntropy supports both class indices and class
         // probabilities as targets.
 
@@ -20,38 +19,38 @@ class LossTests: XCTestCase {
         let indices = MLXArray([0, 1])
         let expected = MLXArray(converting: [0.0, 0.0])
         let loss = crossEntropy(logits: logits, targets: indices, reduction: .none)
-        XCTAssertTrue(allClose(loss, expected).item())
+        expectEqual(loss, expected)
 
         // test class probabilities
         let probs = zeros([2, 2])
         probs[0, 0] = MLXArray(1)
         probs[1, 1] = MLXArray(1)
         let loss2 = crossEntropy(logits: logits, targets: probs, reduction: .none)
-        XCTAssertTrue(all(isNaN(loss2)).item())
+        #expect(all(isNaN(loss2)).item(Bool.self))
     }
 
-    func testCrossEntropyAxis() {
+    @Test func crossEntropyAxis() {
         // The class axis need not be the last one. Expected values from Python mlx:
         // nn.losses.cross_entropy(logits, targets, axis=0) -> [0.169846, 0.169846].
         let logits = MLXArray(converting: [2.0, -1.0, -1.0, 2.0, 0.0, 0.0], [3, 2])
         let targets = MLXArray([0, 1])
         let expected = MLXArray(converting: [0.169846, 0.169846])
         let loss = crossEntropy(logits: logits, targets: targets, axis: 0, reduction: .none)
-        XCTAssertEqual(loss.shape, [2])
-        XCTAssertTrue(allClose(loss, expected, atol: 1e-5).item())
+        #expect(loss.shape == [2])
+        expectEqual(loss, expected, atol: 1e-5)
     }
 
-    func testNLLLossAxis() {
+    @Test func nllLossAxis() {
         // nn.losses.nll_loss(inputs, targets, axis=0) -> [-2, -2]
         let inputs = MLXArray(converting: [2.0, -1.0, -1.0, 2.0, 0.0, 0.0], [3, 2])
         let targets = MLXArray([0, 1])
         let expected = MLXArray(converting: [-2.0, -2.0])
         let loss = nllLoss(inputs: inputs, targets: targets, axis: 0, reduction: .none)
-        XCTAssertEqual(loss.shape, [2])
-        XCTAssertTrue(allClose(loss, expected).item())
+        #expect(loss.shape == [2])
+        expectEqual(loss, expected)
     }
 
-    func testTripletLoss() {
+    @Test func tripletLossReductionsAndNorms() {
         // This mainly tests that tripletLoss now works for arbitrary values of p
 
         let anchors = MLXArray([1, 2, 3, 1, 2, 3], [2, 3])
@@ -62,19 +61,19 @@ class LossTests: XCTestCase {
         let lossesNone = tripletLoss(
             anchors: anchors, positives: positives, negatives: negatives, reduction: .none)
         let expectedNone = MLXArray(converting: [0.0, 2.31662])
-        XCTAssertTrue(allClose(lossesNone, expectedNone).item())
+        expectEqual(lossesNone, expectedNone)
 
         // Test with reduction 'mean'
         let lossesMean = tripletLoss(
             anchors: anchors, positives: positives, negatives: negatives, reduction: .mean)
         let expectedMean = mean(expectedNone)
-        XCTAssertTrue(allClose(lossesMean, expectedMean).item())
+        expectEqual(lossesMean, expectedMean)
 
         // Test with reduction 'sum'
         let lossesSum = tripletLoss(
             anchors: anchors, positives: positives, negatives: negatives, reduction: .sum)
         let expectedSum = sum(expectedNone)
-        XCTAssertTrue(allClose(lossesSum, expectedSum).item())
+        expectEqual(lossesSum, expectedSum)
 
         // Test with non-default 'p' norm degrees
         let normAnchors = MLXArray(converting: [1.0, 2.0]).reshaped([1, 2])
@@ -85,12 +84,12 @@ class LossTests: XCTestCase {
             anchors: normAnchors, positives: normPositives, negatives: normNegatives,
             p: 1, reduction: .none)
         let expectedP1 = MLXArray(converting: [6.0])
-        XCTAssertTrue(allClose(lossesP1, expectedP1).item())
+        expectEqual(lossesP1, expectedP1)
 
         let lossesP3 = tripletLoss(
             anchors: normAnchors, positives: normPositives, negatives: normNegatives,
             p: 3, reduction: .none)
         let expectedP3 = MLXArray(converting: [4.23802])
-        XCTAssertTrue(allClose(lossesP3, expectedP3).item())
+        expectEqual(lossesP3, expectedP3)
     }
 }
diff --git a/Tests/MLXTests/MLXArray+IndexingTests.swift b/Tests/MLXTests/MLXArray+IndexingTests.swift
index 699183f4..3dc1aecd 100644
--- a/Tests/MLXTests/MLXArray+IndexingTests.swift
+++ b/Tests/MLXTests/MLXArray+IndexingTests.swift
@@ -1,7 +1,7 @@
 // Copyright © 2024 Apple Inc.
 
 import Foundation
-import XCTest
+import Testing
 
 @testable import MLX
 
@@ -20,93 +20,91 @@ extension MLXArrayIndexOperation: Equatable {
     }
 }
 
-class MLXArrayIndexingTests: XCTestCase {
-
-    override class func setUp() {
-    }
+@Suite
+struct MLXArrayIndexingTests {
 
     // MARK: - Subscript (get)
 
-    func testArraySubscriptInt() {
+    @Test func arraySubscriptInt() {
         let a = MLXArray(0 ..< 512, [8, 8, 8])
         let s = a[1]
-        XCTAssertEqual(s.ndim, 2)
-        XCTAssertEqual(s.shape, [8, 8])
+        #expect(s.ndim == 2)
+        #expect(s.shape == [8, 8])
 
-        assertEqual(s, MLXArray(64 ..< 128, [8, 8]))
+        expectEqual(s, MLXArray(64 ..< 128, [8, 8]))
     }
 
-    func testArraySubscriptIntArray() {
+    @Test func arraySubscriptIntArray() {
         // squeeze output dimensions as needed
         let a = MLXArray(0 ..< 512, [8, 8, 8])
         let s1 = a[1, 2]
-        XCTAssertEqual(s1.ndim, 1)
-        XCTAssertEqual(s1.shape, [8])
-        assertEqual(s1, MLXArray(80 ..< 88))
+        #expect(s1.ndim == 1)
+        #expect(s1.shape == [8])
+        expectEqual(s1, MLXArray(80 ..< 88))
 
         let s2 = a[1, 2, 3]
-        XCTAssertEqual(s2.ndim, 0)
-        XCTAssertEqual(s2.shape, [])
-        XCTAssertEqual(s2.item(), 64 + 2 * 8 + 3)
+        #expect(s2.ndim == 0)
+        #expect(s2.shape == [])
+        #expect(s2.item(Int.self) == 64 + 2 * 8 + 3)
     }
 
-    func testArraySubscriptIntArray2() {
+    @Test func arraySubscriptIntArray2() {
         // last dimension should not be squeezed
         let a = MLXArray(0 ..< 512, [8, 8, 8, 1])
 
         let s = a[1]
-        XCTAssertEqual(s.ndim, 3)
-        XCTAssertEqual(s.shape, [8, 8, 1])
+        #expect(s.ndim == 3)
+        #expect(s.shape == [8, 8, 1])
 
         let s1 = a[1, 2]
-        XCTAssertEqual(s1.ndim, 2)
-        XCTAssertEqual(s1.shape, [8, 1])
+        #expect(s1.ndim == 2)
+        #expect(s1.shape == [8, 1])
 
         let s2 = a[1, 2, 3]
-        XCTAssertEqual(s2.ndim, 1)
-        XCTAssertEqual(s2.shape, [1])
+        #expect(s2.ndim == 1)
+        #expect(s2.shape == [1])
     }
 
-    func testArraySubscriptFromEnd() {
+    @Test func arraySubscriptFromEnd() {
         // allow negative indices to indicate distance from end (only for int indices)
         let a = MLXArray(0 ..< 12, [3, 4])
         let s = a[-1, -2]
-        XCTAssertEqual(s.ndim, 0)
-        XCTAssertEqual(s.item(), 10)
+        #expect(s.ndim == 0)
+        #expect(s.item(Int.self) == 10)
     }
 
-    func testArraySubscriptRange() {
+    @Test func arraySubscriptRange() {
         let a = MLXArray(0 ..< 512, [8, 8, 8])
 
         let s1 = a[1 ..< 3]
-        XCTAssertEqual(s1.ndim, 3)
-        XCTAssertEqual(s1.shape, [2, 8, 8])
-        assertEqual(s1, MLXArray(64 ..< 192, [2, 8, 8]))
+        #expect(s1.ndim == 3)
+        #expect(s1.shape == [2, 8, 8])
+        expectEqual(s1, MLXArray(64 ..< 192, [2, 8, 8]))
 
         // even though the first dimension is 1 we do not squeeze it
         let s2 = a[1 ... 1]
-        XCTAssertEqual(s2.ndim, 3)
-        XCTAssertEqual(s2.shape, [1, 8, 8])
-        assertEqual(s2, MLXArray(64 ..< 128, [1, 8, 8]))
+        #expect(s2.ndim == 3)
+        #expect(s2.shape == [1, 8, 8])
+        expectEqual(s2, MLXArray(64 ..< 128, [1, 8, 8]))
 
         // multiple ranges, resolving RangeExpressions vs the dimensions
         let s3 = a[1 ..< 2, ..<3, 3...]
-        XCTAssertEqual(s3.ndim, 3)
-        XCTAssertEqual(s3.shape, [1, 3, 5])
-        assertEqual(
+        #expect(s3.ndim == 3)
+        #expect(s3.shape == [1, 3, 5])
+        expectEqual(
             s3, MLXArray([67, 68, 69, 70, 71, 75, 76, 77, 78, 79, 83, 84, 85, 86, 87], [1, 3, 5]))
 
         let s4 = a[-2 ..< -1, ..<(-3), (-3)...]
-        XCTAssertEqual(s4.ndim, 3)
-        XCTAssertEqual(s4.shape, [1, 5, 3])
-        assertEqual(
+        #expect(s4.ndim == 3)
+        #expect(s4.shape == [1, 5, 3])
+        expectEqual(
             s4,
             MLXArray(
                 [389, 390, 391, 397, 398, 399, 405, 406, 407, 413, 414, 415, 421, 422, 423],
                 [1, 5, 3]))
     }
 
-    func testArraySubscriptAdvanced() {
+    @Test func arraySubscriptAdvanced() {
         // advanced subscript examples taken from
         // https://numpy.org/doc/stable/user/basics.indexing.html#integer-array-indexing
 
@@ -117,24 +115,24 @@ class MLXArrayIndexingTests: XCTestCase {
 
         let s1 = a[i1, i2]
 
-        XCTAssertEqual(s1.ndim, 1)
-        XCTAssertEqual(s1.shape, [3])
+        #expect(s1.ndim == 1)
+        #expect(s1.shape == [3])
 
         let expected = MLXArray([0, 15, 30].asInt32)
-        assertEqual(s1, expected)
+        expectEqual(s1, expected)
     }
 
-    func testArraySubscriptAdvanced2() {
+    @Test func arraySubscriptAdvanced2() {
         let a = MLXArray(0 ..< 12, [6, 2]).asType(Int32.self)
 
         let i1 = MLXArray([0, 2, 4])
         let s2 = a[i1]
 
         let expected = MLXArray([0, 1, 4, 5, 8, 9].asInt32, [3, 2])
-        assertEqual(s2, expected)
+        expectEqual(s2, expected)
     }
 
-    func testArraySubscriptAdvanced2d() {
+    @Test func arraySubscriptAdvanced2d() {
         let a = MLXArray(0 ..< 12, [4, 3]).asType(Int32.self)
 
         let rows = MLXArray([0, 0, 3, 3], [2, 2])
@@ -143,10 +141,10 @@ class MLXArrayIndexingTests: XCTestCase {
         let s = a[rows, cols]
 
         let expected = MLXArray([0, 2, 9, 11].asInt32, [2, 2])
-        assertEqual(s, expected)
+        expectEqual(s, expected)
     }
 
-    func testArraySubscriptAdvanced2d2() {
+    @Test func arraySubscriptAdvanced2d2() {
         let a = MLXArray(0 ..< 12, [4, 3]).asType(Int32.self)
 
         let rows = MLXArray([0, 3], [2, 1])
@@ -155,20 +153,20 @@ class MLXArrayIndexingTests: XCTestCase {
         let s = a[rows, cols]
 
         let expected = MLXArray([0, 2, 9, 11].asInt32, [2, 2])
-        assertEqual(s, expected)
+        expectEqual(s, expected)
     }
 
     // MARK: - Subscript (set)
 
-    func testArrayMutateSingleIndex() {
+    @Test func arrayMutateSingleIndex() {
         let a = MLXArray((0 ..< 12).asInt32, [3, 4])
         a[1] = [77]
 
         let expected = MLXArray([0, 1, 2, 3, 77, 77, 77, 77, 8, 9, 10, 11].asInt32, [3, 4])
-        assertEqual(a, expected)
+        expectEqual(a, expected)
     }
 
-    func testArrayMutateBroadcastMultiIndex() {
+    @Test func arrayMutateBroadcastMultiIndex() {
         let a = MLXArray((0 ..< 20).asInt32, [2, 2, 5])
 
         // broadcast to a row
@@ -183,10 +181,10 @@ class MLXArrayIndexingTests: XCTestCase {
         let expected = MLXArray(
             [55, 66, 77, 88, 99, 5, 6, 7, 123, 9, 77, 77, 77, 77, 77, 15, 16, 17, 18, 19].asInt32,
             [2, 2, 5])
-        assertEqual(a, expected)
+        expectEqual(a, expected)
     }
 
-    func testArrayMutateBroadcastSlice() {
+    @Test func arrayMutateBroadcastSlice() {
         let a = MLXArray((0 ..< 20).asInt32, [2, 2, 5])
 
         // write using slices -- this ends up covering two elements
@@ -195,10 +193,10 @@ class MLXArrayIndexingTests: XCTestCase {
         let expected = MLXArray(
             [0, 1, 2, 3, 4, 5, 6, 88, 88, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19].asInt32,
             [2, 2, 5])
-        assertEqual(a, expected)
+        expectEqual(a, expected)
     }
 
-    func testArrayMutateAdvanced() {
+    @Test func arrayMutateAdvanced() {
         let a = MLXArray(0 ..< 35, [5, 7]).asType(Int32.self)
 
         let i1 = MLXArray([0, 2, 4])
@@ -206,18 +204,18 @@ class MLXArrayIndexingTests: XCTestCase {
 
         a[i1, i2] = MLXArray([100, 200, 300].asInt32)
 
-        XCTAssertEqual(a[0, 0].item(), 100.int32)
-        XCTAssertEqual(a[2, 1].item(), 200.int32)
-        XCTAssertEqual(a[4, 2].item(), 300.int32)
+        #expect(a[0, 0].item(Int32.self) == 100)
+        #expect(a[2, 1].item(Int32.self) == 200)
+        #expect(a[4, 2].item(Int32.self) == 300)
     }
 
-    func testCollection() {
+    @Test func collection() {
         let a = MLXArray((0 ..< 20).asInt32, [2, 2, 5])
 
         // enumerate "rows"
         for (i, r) in a.enumerated() {
             let expected = MLXArray((i * 10 ..< i * 10 + 10).asInt32, [2, 5])
-            assertEqual(r, expected)
+            expectEqual(r, expected)
         }
     }
 
@@ -226,11 +224,11 @@ class MLXArrayIndexingTests: XCTestCase {
     // These calls are still tested but have been deprecated in favor of the full python indexing,
     // e.g. array[.ellipsis, .stride(by: 2)]
 
-    func testArraySubscriptIntAxis() {
+    @Test func arraySubscriptIntAxis() {
         let a = MLXArray(0 ..< 512, [8, 8, 8])
         let s = a[1, axis: -1]
-        XCTAssertEqual(s.ndim, 2)
-        XCTAssertEqual(s.shape, [8, 8])
+        #expect(s.ndim == 2)
+        #expect(s.shape == [8, 8])
 
         // array([[1, 9, 17, ..., 41, 49, 57],
         //        [65, 73, 81, ..., 105, 113, 121],
@@ -240,81 +238,81 @@ class MLXArrayIndexingTests: XCTestCase {
         //        [385, 393, 401, ..., 425, 433, 441],
         //        [449, 457, 465, ..., 489, 497, 505]], dtype=int64)
 
-        XCTAssertEqual(s[0, 0].item(Int.self), 1)
-        XCTAssertEqual(s[0, 1].item(Int.self), 9)
-        XCTAssertEqual(s[0, 2].item(Int.self), 17)
+        #expect(s[0, 0].item(Int.self) == 1)
+        #expect(s[0, 1].item(Int.self) == 9)
+        #expect(s[0, 2].item(Int.self) == 17)
     }
 
-    public func testReversed() {
+    @Test func reversed() {
         // tests for [::-1]
         let a = MLXArray(0 ..< 4)
 
         // reverse last dimension (only in this case)
-        assertEqual(a[stride: -1], MLXArray([3, 2, 1, 0]))
+        expectEqual(a[stride: -1], MLXArray([3, 2, 1, 0]))
 
         // reverse first dimension
-        assertEqual(a.reshaped(1, 1, 4)[stride: -1, axis: 0], MLXArray(0 ..< 4).reshaped(1, 1, 4))
+        expectEqual(a.reshaped(1, 1, 4)[stride: -1, axis: 0], MLXArray(0 ..< 4).reshaped(1, 1, 4))
 
     }
 
-    public func testStridedBy2() {
+    @Test func stridedBy2() {
         let a = MLXArray(0 ..< (2 * 3 * 4), [2, 3, 4])
 
         let r = a[stride: 2, axis: -1]
         let expected = MLXArray(Array(stride(from: 0, to: 2 * 3 * 4, by: 2)), [2, 3, 2])
-        assertEqual(r, expected)
+        expectEqual(r, expected)
     }
 
-    public func testStridedBy2Offset() {
+    @Test func stridedBy2Offset() {
         let a = MLXArray(0 ..< (2 * 3 * 5), [2, 3, 5])
 
         let r = a[from: 1, stride: 2, axis: -1]
         let expected = MLXArray([1, 3, 6, 8, 11, 13, 16, 18, 21, 23, 26, 28], [2, 3, 2])
-        assertEqual(r, expected)
+        expectEqual(r, expected)
     }
 
-    public func testStridedByNegative1Last() {
+    @Test func stridedByNegative1Last() {
         let a = MLXArray(0 ..< 6, [2, 3])
 
         let r = a[stride: -1, axis: -1]
         let expected = MLXArray([2, 1, 0, 5, 4, 3], [2, 3])
-        assertEqual(r, expected)
+        expectEqual(r, expected)
     }
 
-    public func testStridedByNegative1First() {
+    @Test func stridedByNegative1First() {
         let a = MLXArray(0 ..< 6, [2, 3])
 
         let r = a[stride: -1, axis: 0]
         let expected = MLXArray([3, 4, 5, 0, 1, 2], [2, 3])
-        assertEqual(r, expected)
+        expectEqual(r, expected)
     }
 
-    public func testStridedByNegative1SecondOffset() {
+    @Test func stridedByNegative1SecondOffset() {
         let a = MLXArray(0 ..< (2 * 3 * 5), [2, 3, 5])
 
         let r = a[from: 1, stride: -1, axis: 1]
         let expected = MLXArray(
             [5, 6, 7, 8, 9, 0, 1, 2, 3, 4, 20, 21, 22, 23, 24, 15, 16, 17, 18, 19], [2, 2, 5])
-        assertEqual(r, expected)
+        expectEqual(r, expected)
     }
 
-    public func testStridedByNegative1LastOffset() {
+    @Test func stridedByNegative1LastOffset() {
         let a = MLXArray(0 ..< (2 * 3 * 5), [2, 3, 5])
 
         let r = a[from: 1, stride: -1, axis: -1]
         let expected = MLXArray([1, 0, 6, 5, 11, 10, 16, 15, 21, 20, 26, 25], [2, 3, 2])
-        assertEqual(r, expected)
+        expectEqual(r, expected)
     }
 
-    public func testStridedByNegative2SecondOffset() {
+    @Test func stridedByNegative2SecondOffset() {
         let a = MLXArray(0 ..< (2 * 5 * 3), [2, 5, 3])
 
         let r = a[from: 1, stride: -2, axis: 1]
         let expected = MLXArray([3, 4, 5, 18, 19, 20], [2, 1, 3])
-        assertEqual(r, expected)
+        expectEqual(r, expected)
     }
 
-    public func testStridedByNegative2Last() {
+    @Test func stridedByNegative2Last() {
         let a = MLXArray(0 ..< (2 * 3 * 4), [2, 3, 4])
 
         let r = a[stride: -2, axis: -1]
@@ -329,10 +327,10 @@ class MLXArrayIndexingTests: XCTestCase {
         //         [23, 21]]], dtype=int64)
 
         let expected = MLXArray([3, 1, 7, 5, 11, 9, 15, 13, 19, 17, 23, 21], [2, 3, 2])
-        assertEqual(r, expected)
+        expectEqual(r, expected)
     }
 
-    public func testStridedByNegative2First() {
+    @Test func stridedByNegative2First() {
         let a = MLXArray(0 ..< (2 * 3 * 4), [2, 3, 4])
 
         let r = a[stride: -2, axis: 0]
@@ -344,10 +342,10 @@ class MLXArrayIndexingTests: XCTestCase {
         //         [20, 21, 22, 23]]], dtype=int64)
 
         let expected = MLXArray([12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23], [1, 3, 4])
-        assertEqual(r, expected)
+        expectEqual(r, expected)
     }
 
-    public func testStridedByNegative2SecondSet() {
+    @Test func stridedByNegative2SecondSet() {
         let a = MLXArray(Int32(0) ..< (2 * 3 * 4), [2, 3, 4])
 
         a[stride: -2, axis: 1] = [99, 88, 77, 66]
@@ -364,20 +362,19 @@ class MLXArrayIndexingTests: XCTestCase {
                 99, 88, 77, 66, 4, 5, 6, 7, 99, 88, 77, 66, 99, 88, 77, 66, 16, 17, 18, 19, 99, 88,
                 77, 66,
             ], [2, 3, 4])
-        assertEqual(a, expected)
+        expectEqual(a, expected)
     }
 
     // MARK: - Full Indexing - Read
 
-    public func testExpandEllipsisOperations() {
+    @Test func expandEllipsis() {
         // test equivalent of mlx_expand_ellipsis
         let shape: [Int32] = [4, 6, 8, 8, 2]
 
         let operations1: [MLXArrayIndexOperation] = [.ellipsis, .index(0), .index(0)]
         let result1 = expandEllipsisOperations(shape: shape, operations: operations1)
-        XCTAssertEqual(
-            result1,
-            [
+        #expect(
+            result1 == [
                 .slice(.init(start: 0, end: 4, stride: 1)),
                 .slice(.init(start: 0, end: 6, stride: 1)),
                 .slice(.init(start: 0, end: 8, stride: 1)),
@@ -387,9 +384,8 @@ class MLXArrayIndexingTests: XCTestCase {
 
         let operations2: [MLXArrayIndexOperation] = [.index(0), .index(0), .ellipsis]
         let result2 = expandEllipsisOperations(shape: shape, operations: operations2)
-        XCTAssertEqual(
-            result2,
-            [
+        #expect(
+            result2 == [
                 .index(0),
                 .index(0),
                 .slice(.init(start: 0, end: 8, stride: 1)),
@@ -399,9 +395,8 @@ class MLXArrayIndexingTests: XCTestCase {
 
         let operations3: [MLXArrayIndexOperation] = [.index(0), .ellipsis, .index(0)]
         let result3 = expandEllipsisOperations(shape: shape, operations: operations3)
-        XCTAssertEqual(
-            result3,
-            [
+        #expect(
+            result3 == [
                 .index(0),
                 .slice(.init(start: 0, end: 6, stride: 1)),
                 .slice(.init(start: 0, end: 8, stride: 1)),
@@ -413,9 +408,8 @@ class MLXArrayIndexingTests: XCTestCase {
             .newAxis, .index(0), .ellipsis, .index(0), .newAxis,
         ]
         let result4 = expandEllipsisOperations(shape: shape, operations: operations4)
-        XCTAssertEqual(
-            result4,
-            [
+        #expect(
+            result4 == [
                 .newAxis,
                 .index(0),
                 .slice(.init(start: 0, end: 6, stride: 1)),
@@ -427,14 +421,14 @@ class MLXArrayIndexingTests: XCTestCase {
     }
 
     func check(
-        _ result: MLXArray, _ shape: [Int], _ sum: Int, file: StaticString = #filePath,
-        line: UInt = #line
+        _ result: MLXArray, _ shape: [Int], _ sum: Int,
+        sourceLocation: SourceLocation = #_sourceLocation
     ) {
-        XCTAssertEqual(result.shape, shape, file: file, line: line)
-        XCTAssertEqual(result.sum().item(Int.self), sum, file: file, line: line)
+        #expect(result.shape == shape, sourceLocation: sourceLocation)
+        #expect(result.sum().item(Int.self) == sum, sourceLocation: sourceLocation)
     }
 
-    public func testFullIndexReadSingle() {
+    @Test func fullIndexReadSingle() {
         // single operators go through an optimized path
 
         // a = mx.arange(60).reshape(3, 4, 5)
@@ -459,7 +453,7 @@ class MLXArrayIndexingTests: XCTestCase {
         check(a[i], [2, 4, 5], 1580)
     }
 
-    public func testFullIndexReadNoArray() {
+    @Test func fullIndexReadNoArray() {
         // a = mx.arange(360).reshape(2, 3, 4, 5, 3)
         let a = MLXArray(0 ..< 360, [2, 3, 4, 5, 3])
 
@@ -487,7 +481,7 @@ class MLXArrayIndexingTests: XCTestCase {
             [2, 2, 2, 5, 1, 2], 13160)
     }
 
-    public func testFullIndexReadArray() {
+    @Test func fullIndexReadArray() {
         // these have an MLXArray as a source of indexes and go through the gather path
 
         // a = mx.arange(540).reshape(3, 3, 4, 5, 3)
@@ -527,18 +521,18 @@ class MLXArrayIndexingTests: XCTestCase {
 
     // MARK: - Full Indexing - Write
 
-    public func testFullIndexWriteSingle() {
+    @Test func fullIndexWriteSingle() {
         // single operators go through an optimized path
 
         func check(
-            _ indexes: [MLXArrayIndex], _ sum: Int, file: StaticString = #filePath,
-            line: UInt = #line
+            _ indexes: [MLXArrayIndex], _ sum: Int,
+            sourceLocation: SourceLocation = #_sourceLocation
         ) {
             // a = mx.arange(60).reshape(3, 4, 5)
             let a = MLXArray(0 ..< 60, [3, 4, 5])
 
             a[operations: indexes.map { $0.mlxArrayIndexOperation }] = MLXArray(1)
-            XCTAssertEqual(a.sum().item(Int.self), sum, file: file, line: line)
+            #expect(a.sum().item(Int.self) == sum, sourceLocation: sourceLocation)
         }
 
         // a[...]
@@ -560,16 +554,16 @@ class MLXArrayIndexingTests: XCTestCase {
         check([i], 230)
     }
 
-    public func testFullIndexWriteNoArray() {
+    @Test func fullIndexWriteNoArray() {
         func check(
-            _ indexes: [MLXArrayIndex], _ sum: Int, file: StaticString = #filePath,
-            line: UInt = #line
+            _ indexes: [MLXArrayIndex], _ sum: Int,
+            sourceLocation: SourceLocation = #_sourceLocation
         ) {
             // a = mx.arange(360).reshape(2, 3, 4, 5, 3)
             let a = MLXArray(0 ..< 360, [2, 3, 4, 5, 3])
 
             a[operations: indexes.map { $0.mlxArrayIndexOperation }] = MLXArray(1)
-            XCTAssertEqual(a.sum().item(Int.self), sum, file: file, line: line)
+            #expect(a.sum().item(Int.self) == sum, sourceLocation: sourceLocation)
         }
 
         // a[..., 0] = 1
@@ -596,18 +590,18 @@ class MLXArrayIndexingTests: XCTestCase {
             51540)
     }
 
-    public func testFullIndexWriteArray() {
+    @Test func fullIndexWriteArray() {
         // these have an MLXArray as a source of indexes and go through the gather path
 
         func check(
-            _ indexes: [MLXArrayIndex], _ sum: Int, file: StaticString = #filePath,
-            line: UInt = #line
+            _ indexes: [MLXArrayIndex], _ sum: Int,
+            sourceLocation: SourceLocation = #_sourceLocation
         ) {
             // a = mx.arange(540).reshape(3, 3, 4, 5, 3)
             let a = MLXArray(0 ..< 540, [3, 3, 4, 5, 3])
 
             a[operations: indexes.map { $0.mlxArrayIndexOperation }] = MLXArray(1)
-            XCTAssertEqual(a.sum().item(Int.self), sum, file: file, line: line)
+            #expect(a.sum().item(Int.self) == sum, sourceLocation: sourceLocation)
         }
 
         // i = mx.array([2, 1])
@@ -639,7 +633,7 @@ class MLXArrayIndexingTests: XCTestCase {
             [.stride(by: -1), ..<2, i, 2..., .ellipsis, .newAxis, .stride(by: 2)], 128142)
     }
 
-    public func testSliceWithBroadcast() {
+    @Test func sliceWithBroadcast() {
         // https://github.com/ml-explore/mlx-swift/issues/76
 
         let a = MLXArray.ones([2, 6, 6, 6])
@@ -658,10 +652,10 @@ class MLXArrayIndexingTests: XCTestCase {
                 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
             ], [3, 4, 4, 4])
 
-        assertEqual(b, expected)
+        expectEqual(b, expected)
     }
 
-    public func testCopyEllipsis() {
+    @Test func copyEllipsis() {
         // https://github.com/ml-explore/mlx-swift/issues/364
         let x = arange(0, 12).reshaped([3, 4])
 
@@ -672,7 +666,7 @@ class MLXArrayIndexingTests: XCTestCase {
         y[.ellipsis, 1] *= 10
 
         // x and y should be different
-        XCTAssertFalse(x.allClose(y).all().item())
+        expectNotEqual(x, y)
     }
 
 }
diff --git a/Tests/MLXTests/MLXArray+InitTests.swift b/Tests/MLXTests/MLXArray+InitTests.swift
index 35e34471..1dab796c 100644
--- a/Tests/MLXTests/MLXArray+InitTests.swift
+++ b/Tests/MLXTests/MLXArray+InitTests.swift
@@ -2,119 +2,114 @@
 
 import Foundation
 import Numerics
-import XCTest
+import Testing
 
 @testable import MLX
 
 #if canImport(IOSurface)
     import IOSurface
+    import CoreVideo
 #endif
 
-class MLXArrayInitTests: XCTestCase {
-
-    override class func setUp() {
-    }
+@Suite
+struct MLXArrayInitTests {
 
     // MARK: - Dtype
-    func testDtypeSize() {
+    @Test func dtypeSize() {
         // Checking that the size of the dtype matches the array's itemsize
         for dtype in DType.allCases {
-            XCTAssertEqual(MLXArray(Data(), dtype: dtype).itemSize, dtype.size)
+            #expect(MLXArray(Data(), dtype: dtype).itemSize == dtype.size)
         }
     }
 
-    func testDtypeCodable() {
+    @Test func dtypeCodable() throws {
         let encoder = JSONEncoder()
         let decoder = JSONDecoder()
         // Test encoding / decoding round trip
         for dtype in DType.allCases {
-            do {
-                let json: Data = try encoder.encode(dtype)
-                let decoded = try decoder.decode(DType.self, from: json)
-                XCTAssertEqual(decoded, dtype)
-            } catch {
-                XCTFail("Encoding / decoding failed")
-            }
+            let json: Data = try encoder.encode(dtype)
+            let decoded = try decoder.decode(DType.self, from: json)
+            #expect(decoded == dtype)
         }
     }
 
     // MARK: - Creation
-    func testInt() {
+    @Test func int() {
         // array creation with Int -- we want it to produce .int32
         let a1 = MLXArray(500)
-        XCTAssertEqual(a1.dtype, .int32)
+        #expect(a1.dtype == .int32)
 
         // eplicit int64
         let a2 = MLXArray(int64: 500)
-        XCTAssertEqual(a2.dtype, .int64)
+        #expect(a2.dtype == .int64)
 
         let a3 = MLXArray([1, 2, 3])
-        XCTAssertEqual(a3.dtype, .int32)
+        #expect(a3.dtype == .int32)
 
         let a4 = MLXArray(int64: [1, 2, 3])
-        XCTAssertEqual(a4.dtype, .int64)
+        #expect(a4.dtype == .int64)
 
         let a5 = MLXArray(0 ..< 12)
-        XCTAssertEqual(a5.dtype, .int32)
+        #expect(a5.dtype == .int32)
 
         let a6 = MLXArray(int64: 0 ..< 12)
-        XCTAssertEqual(a6.dtype, .int64)
+        #expect(a6.dtype == .int64)
     }
 
-    func testArrayCreationLiteralArray() {
+    @Test func arrayCreationLiteralArray() {
         let a: MLXArray = [20, 30, 40]
-        assertEqual(a, MLXArray([20, 30, 40].asInt32))
+        expectEqual(a, MLXArray([20, 30, 40].asInt32))
     }
 
-    func testArrayCreationDoubleArray() {
+    @Test func arrayCreationDoubleArray() {
         // this transforms the array to [Float] and constructs (as a convenience)
         let a = MLXArray(converting: [0.1, 0.5])
-        XCTAssertEqual(a.dtype, .float32)
-        XCTAssertEqual(a[0].item(Float.self), 0.1, accuracy: 0.01)
+        #expect(a.dtype == .float32)
+        expectEqual(a[0].item(Float.self), 0.1, accuracy: 0.01)
     }
 
-    func testArrayCreationArray1D() {
+    @Test func arrayCreationArray1D() {
         let a = MLXArray([1, 2, 3])
-        XCTAssertEqual(a.dtype, .int32)
-        XCTAssertEqual(a.count, 3)
-        XCTAssertEqual(a.ndim, 1)
-        XCTAssertEqual(a.dim(0), 3)
+        #expect(a.dtype == .int32)
+        #expect(a.count == 3)
+        #expect(a.ndim == 1)
+        #expect(a.dim(0) == 3)
     }
 
-    func testArrayCreationRange() {
+    @Test func arrayCreationRange() {
         let a = MLXArray(0 ..< 12, [3, 4])
-        XCTAssertEqual(a.dtype, .int32)
-        XCTAssertEqual(a.size, 12)
-        XCTAssertEqual(a.ndim, 2)
+        #expect(a.dtype == .int32)
+        #expect(a.size == 12)
+        #expect(a.ndim == 2)
     }
 
     // MARK: - Nested (multi dimensional) arrays
 
-    func testArrayCreationNested2D() {
+    @Test func arrayCreationNested2D() {
         let a = MLXArray([
             [1, 2, 3],
             [4, 5, 6],
         ])
-        XCTAssertEqual(a.dtype, .int32)
-        XCTAssertEqual(a.shape, [2, 3])
-        XCTAssertEqual(a.ndim, 2)
-        XCTAssertEqual(a.size, 6)
+        #expect(a.dtype == .int32)
+        #expect(a.shape == [2, 3])
+        #expect(a.ndim == 2)
+        #expect(a.size == 6)
 
         // count is the size of dimension 0
-        XCTAssertEqual(a.count, 2)
+        #expect(a.count == 2)
 
-        assertEqual(a, MLXArray(1 ... 6, [2, 3]))
+        expectEqual(a, MLXArray(1 ... 6, [2, 3]))
     }
 
-    func testArrayCreationNested3D() {
+    @Test func arrayCreationNested3D() {
         // literal nesting and the flat + shape form must agree
         let a = MLXArray([
             [[0, 1], [2, 3], [4, 5]],
             [[6, 7], [8, 9], [10, 11]],
         ])
-        XCTAssertEqual(a.dtype, .int32)
-        XCTAssertEqual(a.shape, [2, 3, 2])
-        assertEqual(a, MLXArray(0 ..< 12, [2, 3, 2]))
+        #expect(a.dtype == .int32)
+        #expect(a.shape == [2, 3, 2])
+        expectEqual(a, MLXArray(0 ..< 12, [2, 3, 2]))
 
         // also works with a computed (non literal) nested array
         let nested = (0 ..< 2).map { i in
@@ -123,252 +118,251 @@ class MLXArrayInitTests: XCTestCase {
             }
         }
         let b = MLXArray(nested)
-        XCTAssertEqual(b.shape, [2, 3, 4])
-        assertEqual(b, MLXArray(0 ..< 24, [2, 3, 4]))
+        #expect(b.shape == [2, 3, 4])
+        expectEqual(b, MLXArray(0 ..< 24, [2, 3, 4]))
     }
 
-    func testArrayCreationNested4D() {
+    @Test func arrayCreationNested4D() {
         let a = MLXArray([
             [[[0, 1], [2, 3]], [[4, 5], [6, 7]]],
             [[[8, 9], [10, 11]], [[12, 13], [14, 15]]],
         ])
-        XCTAssertEqual(a.shape, [2, 2, 2, 2])
-        XCTAssertEqual(a.ndim, 4)
-        assertEqual(a, MLXArray(0 ..< 16, [2, 2, 2, 2]))
+        #expect(a.shape == [2, 2, 2, 2])
+        #expect(a.ndim == 4)
+        expectEqual(a, MLXArray(0 ..< 16, [2, 2, 2, 2]))
     }
 
-    func testArrayCreationNestedDTypes() {
+    @Test func arrayCreationNestedDTypes() {
         // the dtype of a nested array matches the 1d initializers for the same leaf type
-        XCTAssertEqual(MLXArray([[1, 2], [3, 4]]).dtype, .int32)
-        XCTAssertEqual(MLXArray([[Int32(1), 2], [3, 4]]).dtype, .int32)
-        XCTAssertEqual(MLXArray([[Int16(1), 2], [3, 4]]).dtype, .int16)
-        XCTAssertEqual(MLXArray([[UInt8(1), 2], [3, 4]]).dtype, .uint8)
-        XCTAssertEqual(MLXArray([[Float(1), 2], [3, 4]]).dtype, .float32)
-        XCTAssertEqual(MLXArray([[true, false], [false, true]]).dtype, .bool)
+        #expect(MLXArray([[1, 2], [3, 4]]).dtype == .int32)
+        #expect(MLXArray([[Int32(1), 2], [3, 4]]).dtype == .int32)
+        #expect(MLXArray([[Int16(1), 2], [3, 4]]).dtype == .int16)
+        #expect(MLXArray([[UInt8(1), 2], [3, 4]]).dtype == .uint8)
+        #expect(MLXArray([[Float(1), 2], [3, 4]]).dtype == .float32)
+        #expect(MLXArray([[true, false], [false, true]]).dtype == .bool)
 
         // Int64 is a distinct type from Int and must not take the Int -> .int32 path
-        XCTAssertEqual(MLXArray([[Int64(1), 2], [3, 4]]).dtype, .int64)
+        #expect(MLXArray([[Int64(1), 2], [3, 4]]).dtype == .int64)
 
         // casting the literal as a whole is a different inference path than casting the
         // elements and must stay unambiguous
-        XCTAssertEqual(
-            MLXArray([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]] as [[Float32]]).dtype, .float32)
-        XCTAssertEqual(MLXArray([[[1.0, 2.0]]] as [[[Float32]]]).dtype, .float32)
+        #expect(MLXArray([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]] as [[Float32]]).dtype == .float32)
+        #expect(MLXArray([[[1.0, 2.0]]] as [[[Float32]]]).dtype == .float32)
 
         // as with `MLXArray([Double])` this produces .float64 rather than promoting
         let doubles: [[Double]] = [[1, 2], [3, 4]]
-        XCTAssertEqual(MLXArray(doubles).dtype, .float64)
+        #expect(MLXArray(doubles).dtype == .float64)
 
         let complex: [[Complex<Float>]] = [[Complex(2, 7)], [Complex(3, 8)]]
         let c = MLXArray(complex)
-        XCTAssertEqual(c.dtype, .complex64)
-        XCTAssertEqual(c.shape, [2, 1])
-        assertEqual(c, MLXArray([Complex<Float>(2, 7), Complex<Float>(3, 8)], [2, 1]))
+        #expect(c.dtype == .complex64)
+        #expect(c.shape == [2, 1])
+        expectEqual(c, MLXArray([Complex<Float>(2, 7), Complex<Float>(3, 8)], [2, 1]))
     }
 
-    func testArrayCreationNestedInt64() {
+    @Test func arrayCreationNestedInt64() {
         // a value that does not fit in an Int32 -- the whole point of the int64: variant
         let big = Int(Int32.max) + 10
         let a = MLXArray(int64: [[1, 2, 3], [4, 5, big]])
-        XCTAssertEqual(a.dtype, .int64)
-        XCTAssertEqual(a.shape, [2, 3])
-        assertEqual(a, MLXArray(int64: [1, 2, 3, 4, 5, big], [2, 3]))
+        #expect(a.dtype == .int64)
+        #expect(a.shape == [2, 3])
+        expectEqual(a, MLXArray(int64: [1, 2, 3, 4, 5, big], [2, 3]))
     }
 
-    func testArrayCreationNestedConverting() {
+    @Test func arrayCreationNestedConverting() {
         // non square so a transposed flattening could not pass
         let a = MLXArray(converting: [[0.1, 0.5, 0.9], [1.3, 1.7, 2.1]])
-        XCTAssertEqual(a.dtype, .float32)
-        XCTAssertEqual(a.shape, [2, 3])
-        assertEqual(a, MLXArray(converting: [0.1, 0.5, 0.9, 1.3, 1.7, 2.1], [2, 3]))
+        #expect(a.dtype == .float32)
+        #expect(a.shape == [2, 3])
+        expectEqual(a, MLXArray(converting: [0.1, 0.5, 0.9, 1.3, 1.7, 2.1], [2, 3]))
     }
 
-    func testArrayCreationNestedEmpty() {
+    @Test func arrayCreationNestedEmpty() {
         let a = MLXArray([[Int](), [Int]()])
-        XCTAssertEqual(a.shape, [2, 0])
-        XCTAssertEqual(a.size, 0)
+        #expect(a.shape == [2, 0])
+        #expect(a.size == 0)
 
         let b = MLXArray([[Int]]())
-        XCTAssertEqual(b.shape, [0, 0])
-        XCTAssertEqual(b.size, 0)
+        #expect(b.shape == [0, 0])
+        #expect(b.size == 0)
     }
 
-    func testArrayCreationClosedRange() {
+    @Test func arrayCreationClosedRange() {
         let a = MLXArray(Int16(3) ... Int16(6))
-        XCTAssertEqual(a.dtype, .int16)
-        XCTAssertEqual(a.count, 4)
-        XCTAssertEqual(a.ndim, 1)
+        #expect(a.dtype == .int16)
+        #expect(a.count == 4)
+        #expect(a.ndim == 1)
     }
 
-    func testArrayCreationStride() {
+    @Test func arrayCreationStride() {
         let a = MLXArray(stride(from: Float(0.5), to: Float(1.5), by: Float(0.1)))
-        XCTAssertEqual(a.dtype, .float32)
-        XCTAssertEqual(a.count, 10)
-        XCTAssertEqual(a.ndim, 1)
+        #expect(a.dtype == .float32)
+        #expect(a.count == 10)
+        #expect(a.ndim == 1)
     }
 
-    func testArrayCreationZeros() {
+    @Test func arrayCreationZeros() {
         let a = MLXArray.zeros([2, 4], type: Int.self)
-        XCTAssertEqual(a.dtype, .int64)
-        XCTAssertEqual(a.size, 8)
-        XCTAssertEqual(a.ndim, 2)
+        #expect(a.dtype == .int64)
+        #expect(a.size == 8)
+        #expect(a.ndim == 2)
     }
 
     // MARK: - Arange
 
-    func testArangeIntStop() {
+    @Test func arangeIntStop() {
         // arange(10) -> [0, 1, 2, ..., 9]
         let a = arange(10)
-        XCTAssertEqual(a.dtype, .int32)
-        XCTAssertEqual(a.shape, [10])
-        assertEqual(a, MLXArray(0 ..< 10))
+        #expect(a.dtype == .int32)
+        #expect(a.shape == [10])
+        expectEqual(a, MLXArray(0 ..< 10))
     }
 
-    func testArangeIntStartStop() {
+    @Test func arangeIntStartStop() {
         // arange(2, 10) -> [2, 3, 4, ..., 9]
         let a = arange(2, 10)
-        XCTAssertEqual(a.dtype, .int32)
-        XCTAssertEqual(a.shape, [8])
-        assertEqual(a, MLXArray(2 ..< 10))
+        #expect(a.dtype == .int32)
+        #expect(a.shape == [8])
+        expectEqual(a, MLXArray(2 ..< 10))
     }
 
-    func testArangeIntStep() {
+    @Test func arangeIntStep() {
         // arange(2, 10, step: 2) -> [2, 4, 6, 8]
         let a = arange(2, 10, step: 2)
-        XCTAssertEqual(a.dtype, .int32)
-        XCTAssertEqual(a.shape, [4])
-        assertEqual(a, MLXArray([2, 4, 6, 8].asInt32))
+        #expect(a.dtype == .int32)
+        #expect(a.shape == [4])
+        expectEqual(a, MLXArray([2, 4, 6, 8].asInt32))
     }
 
-    func testArangeIntDtype() {
+    @Test func arangeIntDtype() {
         // arange(10, dtype: .float32) -> [0.0, 1.0, ..., 9.0]
         let a = arange(10, dtype: .float32)
-        XCTAssertEqual(a.dtype, .float32)
-        XCTAssertEqual(a.shape, [10])
-        assertEqual(a, MLXArray((0 ..< 10).map { Float($0) }))
+        #expect(a.dtype == .float32)
+        #expect(a.shape == [10])
+        expectEqual(a, MLXArray((0 ..< 10).map { Float($0) }))
     }
 
-    func testArangeIntStepDtype() {
+    @Test func arangeIntStepDtype() {
         // arange(2, 10, step: 2, dtype: .float32) -> [2.0, 4.0, 6.0, 8.0]
         let a = arange(2, 10, step: 2, dtype: .float32)
-        XCTAssertEqual(a.dtype, .float32)
-        XCTAssertEqual(a.shape, [4])
-        assertEqual(a, MLXArray([2.0, 4.0, 6.0, 8.0] as [Float]))
+        #expect(a.dtype == .float32)
+        #expect(a.shape == [4])
+        expectEqual(a, MLXArray([2.0, 4.0, 6.0, 8.0] as [Float]))
     }
 
-    func testArangeDoubleStop() {
+    @Test func arangeDoubleStop() {
         // arange(5.0) -> [0.0, 1.0, 2.0, 3.0, 4.0]
         let a = arange(5.0)
-        XCTAssertEqual(a.dtype, .float32)
-        XCTAssertEqual(a.shape, [5])
-        assertEqual(a, MLXArray([0.0, 1.0, 2.0, 3.0, 4.0] as [Float]))
+        #expect(a.dtype == .float32)
+        #expect(a.shape == [5])
+        expectEqual(a, MLXArray([0.0, 1.0, 2.0, 3.0, 4.0] as [Float]))
     }
 
-    func testArangeDoubleStartStop() {
+    @Test func arangeDoubleStartStop() {
         // arange(1.0, 5.0) -> [1.0, 2.0, 3.0, 4.0]
         let a = arange(1.0, 5.0)
-        XCTAssertEqual(a.dtype, .float32)
-        XCTAssertEqual(a.shape, [4])
-        assertEqual(a, MLXArray([1.0, 2.0, 3.0, 4.0] as [Float]))
+        #expect(a.dtype == .float32)
+        #expect(a.shape == [4])
+        expectEqual(a, MLXArray([1.0, 2.0, 3.0, 4.0] as [Float]))
     }
 
-    func testArangeDoubleStep() {
+    @Test func arangeDoubleStep() {
         // arange(0.0, 2.0, step: 0.5) -> [0.0, 0.5, 1.0, 1.5]
         let a = arange(0.0, 2.0, step: 0.5)
-        XCTAssertEqual(a.dtype, .float32)
-        XCTAssertEqual(a.shape, [4])
-        assertEqual(a, MLXArray([0.0, 0.5, 1.0, 1.5] as [Float]))
+        #expect(a.dtype == .float32)
+        #expect(a.shape == [4])
+        expectEqual(a, MLXArray([0.0, 0.5, 1.0, 1.5] as [Float]))
     }
 
-    func testArangeStaticMethod() {
+    @Test func arangeStaticMethod() {
         // Test static method versions
         let a = MLXArray.arange(10)
-        XCTAssertEqual(a.dtype, .int32)
-        assertEqual(a, MLXArray(0 ..< 10))
+        #expect(a.dtype == .int32)
+        expectEqual(a, MLXArray(0 ..< 10))
 
         let b = MLXArray.arange(0.0, 3.0, step: 0.5)
-        XCTAssertEqual(b.dtype, .float32)
-        assertEqual(b, MLXArray([0.0, 0.5, 1.0, 1.5, 2.0, 2.5] as [Float]))
+        #expect(b.dtype == .float32)
+        expectEqual(b, MLXArray([0.0, 0.5, 1.0, 1.5, 2.0, 2.5] as [Float]))
     }
 
-    func testArangeEmpty() {
+    @Test func arangeEmpty() {
         // arange(0) -> empty array
         let a = arange(0)
-        XCTAssertEqual(a.shape, [0])
+        #expect(a.shape == [0])
 
         // arange(5, 5) -> empty array
         let b = arange(5, 5)
-        XCTAssertEqual(b.shape, [0])
+        #expect(b.shape == [0])
 
         // arange(10, 5) -> empty array (start > stop with positive step)
         let c = arange(10, 5)
-        XCTAssertEqual(c.shape, [0])
+        #expect(c.shape == [0])
     }
 
-    func testData() {
+    @Test func data() {
         let data = Data([1, 2, 3, 4])
         let a = MLXArray(data, [2, 2], type: UInt8.self)
         let b = MLXArray(data, [2, 2], dtype: DType.uint8)
         let expected = MLXArray(UInt8(1) ... 4, [2, 2])
-        assertEqual(a, expected)
-        assertEqual(b, expected)
+        expectEqual(a, expected)
+        expectEqual(b, expected)
     }
 
-    func testUnsafeRawPointer() {
+    @Test func unsafeRawPointer() {
         let data = Data([1, 2, 3, 4])
         let a = data.withUnsafeBytes { ptr in
             MLXArray(ptr, [2, 2], type: UInt8.self)
         }
         let expected = MLXArray(UInt8(1) ... 4, [2, 2])
-        assertEqual(a, expected)
+        expectEqual(a, expected)
     }
 
-    func testUnsafeBufferPointer() {
+    @Test func unsafeBufferPointer() {
         let values: [UInt16] = [1, 2, 3, 4]
         let a = values.withUnsafeBufferPointer { ptr in
             MLXArray(ptr, [2, 2])
         }
         let expected = MLXArray(UInt16(1) ... 4, [2, 2])
-        assertEqual(a, expected)
+        expectEqual(a, expected)
     }
 
-    func testComplexScalar() {
+    @Test func complexScalar() {
         let c1 = MLXArray(real: 3, imaginary: 4)
-        XCTAssertEqual(c1.realPart().item(), 3)
-        XCTAssertEqual(c1.imaginaryPart().item(), 4)
+        #expect(c1.realPart().item(Float.self) == 3)
+        #expect(c1.imaginaryPart().item(Float.self) == 4)
 
         let c2 = MLXArray(Complex(3, 4))
-        assertEqual(c1, c2)
+        expectEqual(c1, c2)
     }
 
-    func testComplexArray() {
+    @Test func complexArray() {
         let r1 = MLXArray(converting: [2, 3, 4])
         let i1 = MLXArray(converting: [7, 8, 9])
         let c1 = r1 + i1.asImaginary()
 
-        assertEqual(c1.realPart(), r1)
-        assertEqual(c1.imaginaryPart(), i1)
+        expectEqual(c1.realPart(), r1)
+        expectEqual(c1.imaginaryPart(), i1)
 
         let a1: [Complex<Float>] = [Complex(2, 7), Complex(3, 8), Complex(4, 9)]
         let c2 = MLXArray(a1)
 
-        assertEqual(c1, c2)
+        expectEqual(c1, c2)
     }
 
-    func testFloat64Array() {
+    @Test func float64Array() {
         let d: [Double] = [1.0, 2.0, 3.0]
         let a = MLXArray(d)
-        XCTAssertEqual(a.dtype, .float64)
+        #expect(a.dtype == .float64)
 
         let b = MLXArray(0.5)
-        XCTAssertEqual(b.dtype, .float32)
+        #expect(b.dtype == .float32)
 
         let c = MLXArray(1.1e40)
 
-        XCTAssertEqual(c.dtype, .float64)
+        #expect(c.dtype == .float64)
 
         let e = MLXArray(float64: 0.5)
-        XCTAssertEqual(e.dtype, .float64)
+        #expect(e.dtype == .float64)
     }
 
     /// The finalizer's captures must be released once mlx has called it.
@@ -379,10 +373,10 @@ class MLXArrayInitTests: XCTestCase {
     /// consuming the reference, the closure — and everything it captured — is
     /// pinned for the lifetime of the process, once per adopted buffer.
     ///
-    /// `testIOSurface` below documents the intended behaviour in a comment
+    /// `ioSurface()` below documents the intended behaviour in a comment
     /// ("implicitly releases it when it returns") but never asserts it, which is
     /// what let this go unnoticed.
-    func testFinalizerCapturesAreReleasedAfterFinalizerRuns() {
+    @Test func finalizerCapturesAreReleasedAfterFinalizerRuns() {
         // mlx may invoke the dtor from its own thread, so the counters cannot be
         // plain locals read across that boundary.
         final class Counter: @unchecked Sendable {
@@ -422,7 +416,7 @@ class MLXArrayInitTests: XCTestCase {
                 finalizerRuns.increment()
                 buffer.deallocate()
             }
-            XCTAssertEqual(array.sum().item(Float.self), 4)
+            #expect(array.sum().item(Float.self) == 4)
         }
 
         // mlx keeps freed buffers in its allocator cache, so the dtor has not
@@ -431,55 +425,55 @@ class MLXArrayInitTests: XCTestCase {
         let deadline = Date().addingTimeInterval(5)
         while finalizerRuns.current == 0, Date() < deadline { usleep(1_000) }
 
-        XCTAssertEqual(finalizerRuns.current, 1, "mlx must call the finalizer exactly once")
-        XCTAssertEqual(
-            witnessReleases.current, 1,
-            "the finalizer's captures must be released along with it — otherwise every adopted "
-                + "buffer permanently pins whatever the closure held")
+        #expect(finalizerRuns.current == 1, "mlx must call the finalizer exactly once")
+        #expect(
+            witnessReleases.current == 1,
+            "the finalizer's captures must be released along with it — otherwise every adopted buffer permanently pins whatever the closure held"
+        )
     }
 
     // MARK: - linspace
 
-    func testLinspaceIntegerBoundsAreFloat() {
+    @Test func linspaceIntegerBoundsAreFloat() {
         // matches python: mx.linspace(0, 1, 3) -> float32 [0, 0.5, 1].
         // an integer result would truncate to [0, 0, 1]
         let a = MLXArray.linspace(0, 1, count: 3)
-        XCTAssertEqual(a.dtype, .float32)
-        assertEqual(a, MLXArray(converting: [0.0, 0.5, 1.0]), atol: 1e-6)
+        #expect(a.dtype == .float32)
+        expectEqual(a, MLXArray(converting: [0.0, 0.5, 1.0]), atol: 1e-6)
     }
 
-    func testLinspaceIntegerDType() {
+    @Test func linspaceIntegerDType() {
         // an integer (truncating) result is available, but has to be asked for
         let a = MLXArray.linspace(0, 10, count: 6, dtype: .int32)
-        XCTAssertEqual(a.dtype, .int32)
-        XCTAssertEqual(a.asArray(Int32.self), [0, 2, 4, 6, 8, 10])
+        #expect(a.dtype == .int32)
+        #expect(a.asArray(Int32.self) == [0, 2, 4, 6, 8, 10])
     }
 
-    func testLinspaceDoubleBoundsAreFloat32() {
+    @Test func linspaceDoubleBoundsAreFloat32() {
         // Double.dtype is float64, but like MLXArray(1.0) we do not promote to
         // float64 -- it is not available on the GPU
         let a = MLXArray.linspace(0.0, 1.0, count: 5)
-        XCTAssertEqual(a.dtype, .float32)
-        assertEqual(a, MLXArray(converting: [0.0, 0.25, 0.5, 0.75, 1.0]), atol: 1e-6)
+        #expect(a.dtype == .float32)
+        expectEqual(a, MLXArray(converting: [0.0, 0.25, 0.5, 0.75, 1.0]), atol: 1e-6)
     }
 
-    func testLinspaceFloatBoundsKeepTheirDType() {
-        XCTAssertEqual(MLXArray.linspace(Float(0), Float(1), count: 5).dtype, .float32)
+    @Test func linspaceFloatBoundsKeepTheirDType() {
+        #expect(MLXArray.linspace(Float(0), Float(1), count: 5).dtype == .float32)
         #if !arch(x86_64)
-            XCTAssertEqual(MLXArray.linspace(Float16(0), Float16(1), count: 5).dtype, .float16)
+            #expect(MLXArray.linspace(Float16(0), Float16(1), count: 5).dtype == .float16)
         #endif
     }
 
-    func testLinspaceEndpoint() {
+    @Test func linspaceEndpoint() {
         let inclusive = MLX.linspace(0, 1, count: 5)
-        assertEqual(inclusive, MLXArray(converting: [0, 0.25, 0.5, 0.75, 1.0]), atol: 1e-6)
+        expectEqual(inclusive, MLXArray(converting: [0, 0.25, 0.5, 0.75, 1.0]), atol: 1e-6)
 
         let halfOpen = MLX.linspace(0, 1, count: 5, endpoint: false)
-        assertEqual(halfOpen, MLXArray(converting: [0, 0.2, 0.4, 0.6, 0.8]), atol: 1e-6)
+        expectEqual(halfOpen, MLXArray(converting: [0, 0.2, 0.4, 0.6, 0.8]), atol: 1e-6)
     }
 
     #if canImport(IOSurface)
-        func testIOSurface() {
+        @Test func ioSurface() throws {
             let height = 100
             let width = 128
             let pixelFormat = kCVPixelFormatType_32BGRA
@@ -491,10 +485,8 @@ class MLXArrayInitTests: XCTestCase {
                 .bytesPerElement: 4,
             ]
 
-            guard let ioSurface = IOSurface(properties: properties) else {
-                XCTFail("unable to allocate IOSurface")
-                return
-            }
+            let ioSurface = try #require(
+                IOSurface(properties: properties), "unable to allocate IOSurface")
 
             let array = MLXArray(
                 rawPointer: ioSurface.baseAddress, [height, width, 4], dtype: .uint8
diff --git a/Tests/MLXTests/MLXArray+NormlizerTests.swift b/Tests/MLXTests/MLXArray+NormlizerTests.swift
index 8fbdbfd6..17fd9a5b 100644
--- a/Tests/MLXTests/MLXArray+NormlizerTests.swift
+++ b/Tests/MLXTests/MLXArray+NormlizerTests.swift
@@ -1,15 +1,16 @@
 // Copyright © 2026 Apple Inc.
 
 import Foundation
-import XCTest
+import Testing
 
 @testable import MLX
 
-public final class MLXArrayL2NormalizationTests: XCTestCase {
+@Suite
+struct MLXArrayL2NormalizationTests {
 
     /// Tests standard L2 normalization for a 1D vector.
     /// Magnitude is exactly 5.0, result should be unit length (1.0).
-    func testL2NormalizationStandard() {
+    @Test func l2NormalizationStandard() {
         let rawArray: [Float] = [3.0, 4.0]
         let array = MLXArray(rawArray, [2])
         let normalized = array.l2Normalized()
@@ -17,16 +18,16 @@ public final class MLXArrayL2NormalizationTests: XCTestCase {
         let rawExpected: [Float] = [0.6, 0.8]
         let expected = MLXArray(rawExpected, [2])
 
-        // Use allClose for floating point comparison in MLX
-        XCTAssertTrue(allClose(normalized, expected).item(Bool.self))
+        // expectEqual uses allClose (plus a shape check) for floating point comparison
+        expectEqual(normalized, expected)
 
         // Verify Magnitude: Must be 1.0
         let magnitude = MLXLinalg.norm(normalized, ord: 2).item(Float.self)
-        XCTAssertEqual(magnitude, 1.0, accuracy: 1e-6)
+        expectEqual(magnitude, 1.0, accuracy: 1e-6)
     }
 
     /// Tests normalization along a specific axis in a 2D matrix.
-    func testL2NormalizationAlongAxis() {
+    @Test func l2NormalizationAlongAxis() {
         // 2x2 Matrix: [[3, 4], [0, 1]]
         let rawArray: [Float] = [3.0, 4.0, 0.0, 1.0]
         let array = MLXArray(rawArray, [2, 2])
@@ -38,11 +39,11 @@ public final class MLXArrayL2NormalizationTests: XCTestCase {
         let rawExpected: [Float] = [0.6, 0.8, 0.0, 1.0]
         let expected = MLXArray(rawExpected, [2, 2])
 
-        XCTAssertTrue(allClose(normalized, expected).item(Bool.self))
+        expectEqual(normalized, expected)
     }
 
     /// CRITICAL: Tests behavior with a zero vector to ensure numerical stability via epsilon.
-    func testL2NormalizationZeroVector() {
+    @Test func l2NormalizationZeroVector() {
         let eps: Float = 1e-8
         let rawArray: [Float] = [0.0, 0.0]
         let array = MLXArray(rawArray, [2])
@@ -53,16 +54,16 @@ public final class MLXArrayL2NormalizationTests: XCTestCase {
         let rawExpected: [Float] = [0.0, 0.0]
         let expected = MLXArray(rawExpected, [2])
 
-        XCTAssertTrue(allClose(normalized, expected).item(Bool.self))
+        expectEqual(normalized, expected)
 
         // Magnitude should be 0.0, not NaN!
         let magnitude = MLXLinalg.norm(normalized, ord: 2).item(Float.self)
-        XCTAssertFalse(magnitude.isNaN, "Resulting magnitude should not be NaN")
-        XCTAssertEqual(magnitude, 0.0)
+        #expect(!magnitude.isNaN, "Resulting magnitude should not be NaN")
+        #expect(magnitude == 0.0)
     }
 
     /// Tests values that are smaller than the provided epsilon.
-    func testL2NormalizationUnderEpsilon() {
+    @Test func l2NormalizationUnderEpsilon() {
         let eps: Float = 1e-3
         let rawArray: [Float] = [1e-5, 1e-5]
         let array = MLXArray(rawArray, [2])  // Norm is approx 1.41 * 1e-5
@@ -72,7 +73,7 @@ public final class MLXArrayL2NormalizationTests: XCTestCase {
         let expectedValue = Float(1e-5) / eps
         let expected = MLXArray([expectedValue, expectedValue], [2])
 
-        XCTAssertTrue(allClose(normalized, expected).item(Bool.self))
+        expectEqual(normalized, expected)
     }
 
 }
diff --git a/Tests/MLXTests/MLXArray+OpsTests.swift b/Tests/MLXTests/MLXArray+OpsTests.swift
index 684e1039..7458005b 100644
--- a/Tests/MLXTests/MLXArray+OpsTests.swift
+++ b/Tests/MLXTests/MLXArray+OpsTests.swift
@@ -2,18 +2,16 @@
 
 import Foundation
 import Numerics
-import XCTest
+import Testing
 
 @testable import MLX
 
-class MLXArrayOpsTests: XCTestCase {
-
-    override class func setUp() {
-    }
+@Suite
+struct MLXArrayOpsTests {
 
     // MARK: - Operators
 
-    func testArithmeticSimple() {
+    @Test func arithmeticSimple() {
         var a = MLXArray([1, 2, 3])
         let b = MLXArray(converting: [-5.0, 37.5, 4])
 
@@ -21,22 +19,22 @@ class MLXArrayOpsTests: XCTestCase {
         let r = a + b - 1
 
         // make sure everything got hooked up
-        XCTAssertEqual(r.shape, [3])
-        XCTAssertEqual(r.dtype, .float32)
-        assertEqual(r, MLXArray(converting: [-5, 38.5, 6]))
+        #expect(r.shape == [3])
+        #expect(r.dtype == .float32)
+        expectEqual(r, MLXArray(converting: [-5, 38.5, 6]))
 
         a += b
-        XCTAssertEqual(a.shape, [3])
-        XCTAssertEqual(a.dtype, .float32)
-        assertEqual(a, MLXArray(converting: [-4, 39.5, 7]))
+        #expect(a.shape == [3])
+        #expect(a.dtype == .float32)
+        expectEqual(a, MLXArray(converting: [-4, 39.5, 7]))
 
         a += 1
-        XCTAssertEqual(a.shape, [3])
-        XCTAssertEqual(a.dtype, .float32)
-        assertEqual(a, MLXArray(converting: [-3, 40.5, 8]))
+        #expect(a.shape == [3])
+        #expect(a.dtype == .float32)
+        expectEqual(a, MLXArray(converting: [-3, 40.5, 8]))
     }
 
-    func testArithmeticMatrix() {
+    @Test func arithmeticMatrix() {
         let a = MLXArray([1, 2, 3, 4], [2, 2])
         let b = MLXArray(converting: [-5.0, 37.5, 4, 7, 1, 0], [2, 3])
 
@@ -44,143 +42,146 @@ class MLXArrayOpsTests: XCTestCase {
         let r = ((a * 5) ** 2).matmul(b)
 
         // make sure everything got hooked up
-        XCTAssertEqual(r.shape, [2, 3])
-        XCTAssertEqual(r.dtype, .float32)
-        assertEqual(r, MLXArray(converting: [575, 1037.5, 100, 1675, 8837.5, 900], [2, 3]))
+        #expect(r.shape == [2, 3])
+        #expect(r.dtype == .float32)
+        expectEqual(r, MLXArray(converting: [575, 1037.5, 100, 1675, 8837.5, 900], [2, 3]))
     }
 
-    func testArithmeticDivide() {
+    @Test func arithmeticDivide() {
         let a = MLXArray([1, 2, 3])
 
         // in python: (a // 2) / 0.5
         let r = a.floorDivide(2) / Float(0.5)
 
         // make sure everything got hooked up
-        XCTAssertEqual(r.shape, [3])
-        XCTAssertEqual(r.dtype, .float32)
-        assertEqual(r, MLXArray(converting: [0, 2, 2]))
+        #expect(r.shape == [3])
+        #expect(r.dtype == .float32)
+        expectEqual(r, MLXArray(converting: [0, 2, 2]))
     }
 
-    func testOpsLogical() {
+    @Test func opsLogical() {
         let a = MLXArray([1, 2, 3])
 
         let r = (a .< (a + 1))
 
         // make sure everything got hooked up
-        XCTAssertEqual(r.shape, [3])
-        XCTAssertEqual(r.dtype, .bool)
-        XCTAssertEqual(r.all().item(Bool.self), true)
-        XCTAssertEqual(r.all().item(), true)
+        #expect(r.shape == [3])
+        #expect(r.dtype == .bool)
+        #expect(r.all().item(Bool.self) == true)
+
+        // generic item() inferred from context
+        let allTrue: Bool = r.all().item()
+        #expect(allTrue)
     }
 
-    func testOpsLogicalBoolContext() {
+    @Test func opsLogicalBoolContext() {
         let a = MLXArray([1, 2, 3])
 
         if (a .< (a + 1)).all().item() {
             // expected
         } else {
-            XCTFail("should be true")
+            Issue.record("should be true")
         }
     }
 
-    func testFunctions() {
+    @Test func functions() {
         let a = MLXArray(0 ..< 12, [4, 3])
 
         let r = a.square().sqrt().transposed().T
-        assertEqual(a, r)
+        expectEqual(a, r)
     }
 
-    func testSplitEqual() {
+    @Test func splitEqual() {
         let a = MLXArray(0 ..< 12, [4, 3])
 
         let s = a.split(parts: 2)
-        XCTAssertEqual(s[0].shape, [2, 3])
-        XCTAssertEqual(s[1].shape, [2, 3])
+        #expect(s[0].shape == [2, 3])
+        #expect(s[1].shape == [2, 3])
     }
 
-    func testSplitIndices() {
+    @Test func splitIndices() {
         let a = MLXArray(0 ..< 12, [4, 3])
 
         let s = a.split(indices: [0, 3])
-        XCTAssertEqual(s[0].shape, [0, 3])
-        XCTAssertEqual(s[1].shape, [3, 3])
-        XCTAssertEqual(s[2].shape, [1, 3])
+        #expect(s[0].shape == [0, 3])
+        #expect(s[1].shape == [3, 3])
+        #expect(s[2].shape == [1, 3])
     }
 
-    func testFlatten() {
+    @Test func flatten() {
         let a = MLXArray(0 ..< (8 * 4 * 3), [8, 4, 3])
 
         let f1 = a.flattened()
-        XCTAssertEqual(f1.shape, [8 * 4 * 3])
+        #expect(f1.shape == [8 * 4 * 3])
 
         let f2 = a.flattened(start: 1)
-        XCTAssertEqual(f2.shape, [8, 4 * 3])
+        #expect(f2.shape == [8, 4 * 3])
     }
 
-    public func testMoveAxis() {
+    @Test func moveAxis() {
         let array = MLXArray(0 ..< 16, [2, 2, 2, 2])
 
         let r = array.movedAxis(source: 0, destination: 3)
 
         let expected = MLXArray(
             [0, 8, 1, 9, 2, 10, 3, 11, 4, 12, 5, 13, 6, 14, 7, 15], [2, 2, 2, 2])
-        assertEqual(r, expected)
+        expectEqual(r, expected)
     }
 
-    public func testSwapAxes() {
+    @Test func swapAxes() {
         let array = MLXArray(0 ..< 16, [2, 2, 2, 2])
 
         let r = array.swappedAxes(0, 3)
 
         let expected = MLXArray(
             [0, 8, 2, 10, 4, 12, 6, 14, 1, 9, 3, 11, 5, 13, 7, 15], [2, 2, 2, 2])
-        assertEqual(r, expected)
+        expectEqual(r, expected)
     }
 
-    public func testPromotionTests() {
+    @Test func promotionTests() {
         // expect that adding a scalar will not promote to a larger type
         let x = MLXArray.ones([10], type: Float16.self)
 
-        XCTAssertEqual((x * 15.5).dtype, .float16)
-        XCTAssertEqual((15.5 * x).dtype, .float16)
+        #expect((x * 15.5).dtype == .float16)
+        #expect((15.5 * x).dtype == .float16)
     }
 
-    public func testBFloat16() {
+    @Test func bfloat16() {
         let x = MLXArray([1, 5, 9]).asType(.bfloat16)
         let y = x * x
-        XCTAssertEqual(y.dtype, .bfloat16)
+        #expect(y.dtype == .bfloat16)
 
         // this will convert the 3 to bfloat16 via init<T: HasDType>(_ value: T, dtype: DType)
         let z = (y + 3).sum()
-        XCTAssertEqual(z.dtype, .bfloat16)
+        #expect(z.dtype == .bfloat16)
 
         // read out the value as a usable type
         let r1 = z.item(Float32.self)
-        XCTAssertEqual(r1, 116)
+        #expect(r1 == 116)
 
         let r2 = z.asArray(Float.self)
-        XCTAssertEqual(r2, [116])
+        #expect(r2 == [116])
     }
 
-    public func testComplexPromote() {
+    @Test func complexPromote() {
         let x = MLXArray(Complex(3, 4))
         let y = x * x
-        XCTAssertEqual(y.dtype, .complex64)
+        #expect(y.dtype == .complex64)
 
         // this will convert the 3 to (3, 0i) via init<T: HasDType>(_ value: T, dtype: DType)
         let z = y + 3
-        XCTAssertEqual(z.dtype, .complex64)
+        #expect(z.dtype == .complex64)
 
         // read out as complex
         let r1 = z.item(Complex.self)
-        XCTAssertEqual(r1, Complex(-4, 24))
+        #expect(r1 == Complex(-4, 24))
 
         // read out the value as a usable type
         let r2 = z.asArray(Complex.self)
-        XCTAssertEqual(r2, [Complex(-4, 24)])
+        #expect(r2 == [Complex(-4, 24)])
     }
 
-    func testCountNonzero() {
+    @Test func countNonzero() {
         let a = MLXArray([0, 1, 2, 0, 3, 0], [2, 3])
         let resAxes0 = a.countNonzero(axes: [0])
         let resAxes1 = a.countNonzero(axes: [1])
@@ -188,11 +189,11 @@ class MLXArrayOpsTests: XCTestCase {
         let resAxis1 = a.countNonzero(axis: 1)
         let resAll = a.countNonzero()
 
-        assertEqual(resAxes0, [0, 2, 1])
-        assertEqual(resAxes0, resAxis0)
-        assertEqual(resAxes1, [2, 1])
-        assertEqual(resAxes1, resAxis1)
-        assertEqual(resAll, 3.asMLXArray(dtype: .int32))
+        expectEqual(resAxes0, [0, 2, 1])
+        expectEqual(resAxes0, resAxis0)
+        expectEqual(resAxes1, [2, 1])
+        expectEqual(resAxes1, resAxis1)
+        expectEqual(resAll, 3.asMLXArray(dtype: .int32))
     }
 
 }
diff --git a/Tests/MLXTests/MLXArrayTests.swift b/Tests/MLXTests/MLXArrayTests.swift
index 2d0dc427..f89ca16f 100644
--- a/Tests/MLXTests/MLXArrayTests.swift
+++ b/Tests/MLXTests/MLXArrayTests.swift
@@ -1,193 +1,191 @@
 // Copyright © 2024 Apple Inc.
 
 import Foundation
-import XCTest
+import Testing
 
 @testable import MLX
 
-class MLXArrayTests: XCTestCase {
+@Suite
+struct MLXArrayTests {
 
-    override class func setUp() {
-    }
-
-    func testArrayProperties() {
+    @Test func arrayProperties() {
         let a = MLXArray(converting: [3.5, 4.5, 5.5, 7.0, 9.4, 10.0], [2, 3, 1])
 
-        XCTAssertEqual(a.itemSize, 4)
-        XCTAssertEqual(a.size, 6)
-        XCTAssertEqual(a.count, 2)
-        XCTAssertEqual(a.nbytes, 6 * 4)
-        XCTAssertEqual(a.ndim, 3)
-        XCTAssertEqual(a.dtype, .float32)
-        XCTAssertEqual(a.shape, [2, 3, 1])
+        #expect(a.itemSize == 4)
+        #expect(a.size == 6)
+        #expect(a.count == 2)
+        #expect(a.nbytes == 6 * 4)
+        #expect(a.ndim == 3)
+        #expect(a.dtype == .float32)
+        #expect(a.shape == [2, 3, 1])
     }
 
-    func testArrayRead() {
+    @Test func arrayRead() {
         let a = MLXArray(0 ..< 12, [4, 3])
 
-        XCTAssertEqual(a.asArray(Int32.self), Array(0 ..< 12))
-        XCTAssertEqual(a[1][2].item(Int.self), 5)
+        #expect(a.asArray(Int32.self) == Array(0 ..< 12))
+        #expect(a[1][2].item(Int.self) == 5)
     }
 
-    func testAsArrayContiguous() {
+    @Test func asArrayContiguous() {
         // read array from contiguous memory
         let a = MLXArray(0 ..< 12, [4, 3])
         let b = a.asArray(Int.self)
-        XCTAssertEqual(b, [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11])
+        #expect(b == [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11])
     }
 
-    func testAsArrayEmpty() {
+    @Test func asArrayEmpty() {
         // an empty slice has size 0 and no backing pointer; it should return
         // an empty array rather than crashing (see issue #480)
         let a = MLXArray.ones([1, 200])
         let b = a[128...]
-        XCTAssertEqual(b.shape, [0, 200])
-        XCTAssertEqual(b.size, 0)
-        XCTAssertEqual(b.asArray(Int.self), [])
-        XCTAssertEqual(b.asArray(Float.self), [])
+        #expect(b.shape == [0, 200])
+        #expect(b.size == 0)
+        #expect(b.asArray(Int.self) == [])
+        #expect(b.asArray(Float.self) == [])
 
         // a directly-constructed empty array behaves the same
         let empty = MLXArray([Int32](), [0])
-        XCTAssertEqual(empty.asArray(Int32.self), [])
+        #expect(empty.asArray(Int32.self) == [])
 
         // a non-contiguous empty array goes through the strided copy path
         let t = b.T
-        XCTAssertEqual(t.shape, [200, 0])
-        XCTAssertEqual(t.asArray(Float.self), [])
+        #expect(t.shape == [200, 0])
+        #expect(t.asArray(Float.self) == [])
     }
 
-    func testAsDataEmpty() {
+    @Test func asDataEmpty() {
         // asData shares physicalSize with asArray, so it also returns empty
         // data for an empty array (see issue #480)
         let b = MLXArray.ones([1, 200])[128...]
         for access in [MLXArray.AccessMethod.copy, .noCopyIfContiguous, .noCopy] {
             let data = b.asData(access: access)
-            XCTAssertEqual(data.data.count, 0)
-            XCTAssertEqual(data.shape, [0, 200])
+            #expect(data.data.count == 0)
+            #expect(data.shape == [0, 200])
         }
-        XCTAssertEqual(b.T.asData(access: .copy).data.count, 0)
+        #expect(b.T.asData(access: .copy).data.count == 0)
     }
 
-    func testAsArrayScalar() {
+    @Test func asArrayScalar() {
         // a scalar (0-dimensional) array yields a single-element array
         let a = MLXArray(7)
-        XCTAssertEqual(a.shape, [])
-        XCTAssertEqual(a.size, 1)
-        XCTAssertEqual(a.asArray(Int32.self), [7])
+        #expect(a.shape == [])
+        #expect(a.size == 1)
+        #expect(a.asArray(Int32.self) == [7])
     }
 
-    func testAsArrayNonContiguous1() {
+    @Test func asArrayNonContiguous1() {
         // skipping elements via slicing
         let a = MLXArray(0 ..< 9, [3, 3])
 
         let s = a[0 ..< 2, 1 ..< 3]
-        assertEqual(s, MLXArray([1, 2, 4, 5], [2, 2]))
+        expectEqual(s, MLXArray([1, 2, 4, 5], [2, 2]))
 
-        XCTAssertEqual(s.shape, [2, 2])
+        #expect(s.shape == [2, 2])
 
         // size and nbytes are the logical size
-        XCTAssertEqual(s.size, 2 * 2)
-        XCTAssertEqual(s.nbytes, 2 * 2 * s.itemSize)
+        #expect(s.size == 2 * 2)
+        #expect(s.nbytes == 2 * 2 * s.itemSize)
 
         // internal property for counting the physical size of the backing.
         // note that the physical size doesn't include the row that is
         // sliced out
-        XCTAssertEqual(s.physicalSize, 3 * 2)
+        #expect(s.physicalSize == 3 * 2)
 
         // evaluating s (the comparison above) will realize the strides.
         // if we examine these before they might be [2, 1] which are the
         // "logical" strides
-        XCTAssertEqual(s.internalStrides, [3, 1])
+        #expect(s.internalStrides == [3, 1])
 
         let s_arr = s.asArray(Int32.self)
-        XCTAssertEqual(s_arr, [1, 2, 4, 5])
+        #expect(s_arr == [1, 2, 4, 5])
     }
 
-    func testAsArrayNonContiguous2() {
+    @Test func asArrayNonContiguous2() {
         // a transpose via strides
         let a = MLXArray(0 ..< 12, [4, 3])
 
         let s = asStrided(a, [3, 4], strides: [1, 3])
 
         let expected: [Int32] = [0, 3, 6, 9, 1, 4, 7, 10, 2, 5, 8, 11]
-        assertEqual(s, MLXArray(expected, [3, 4]))
+        expectEqual(s, MLXArray(expected, [3, 4]))
 
         // Note: be careful to use the matching type -- if we transcode
         // to a different type it will be converted to contiguous
         let s_arr = s.asArray(Int32.self)
-        XCTAssertEqual(s_arr, expected)
+        #expect(s_arr == expected)
     }
 
-    func testAsArrayNonContiguous4() {
+    @Test func asArrayNonContiguous4() {
         // buffer with holes (last dimension has stride of 2 and
         // thus larger storage than it physically needs)
         let a = MLXArray(0 ..< 16, [4, 4])
         let s = a[0..., .stride(by: 2)]
 
         let expected: [Int32] = [0, 2, 4, 6, 8, 10, 12, 14]
-        assertEqual(s, MLXArray(expected, [4, 2]))
+        expectEqual(s, MLXArray(expected, [4, 2]))
 
-        XCTAssertEqual(s.internalStrides, [4, 2])
+        #expect(s.internalStrides == [4, 2])
 
         let s_arr = s.asArray(Int32.self)
-        XCTAssertEqual(s_arr, expected)
+        #expect(s_arr == expected)
     }
 
-    func testContiguousStrides() {
-        XCTAssertEqual(contiguousStrides(shape: [1, 1, 1]), [1, 1, 1])
-        XCTAssertEqual(contiguousStrides(shape: [4, 4]), [4, 1])
-        XCTAssertEqual(contiguousStrides(shape: [4, 2]), [2, 1])
-        XCTAssertEqual(contiguousStrides(shape: [3, 2, 1, 5]), [10, 5, 5, 1])
+    @Test func contiguousStridesForShapes() {
+        #expect(contiguousStrides(shape: [1, 1, 1]) == [1, 1, 1])
+        #expect(contiguousStrides(shape: [4, 4]) == [4, 1])
+        #expect(contiguousStrides(shape: [4, 2]) == [2, 1])
+        #expect(contiguousStrides(shape: [3, 2, 1, 5]) == [10, 5, 5, 1])
     }
 
-    func testAsDataContiguous() {
+    @Test func asDataContiguous() {
         // contiguous source
         let a = MLXArray(0 ..< 16, [4, 4])
 
         let result = a.asData()
-        XCTAssertEqual(result.shape, [4, 4])
-        XCTAssertEqual(result.strides, [4, 1])
-        XCTAssertEqual(result.dType, .int32)
+        #expect(result.shape == [4, 4])
+        #expect(result.strides == [4, 1])
+        #expect(result.dType == .int32)
     }
 
-    func testAsDataContiguousNoCopy() {
+    @Test func asDataContiguousNoCopy() {
         // contiguous source
         let a = MLXArray(0 ..< 16, [4, 4])
 
         do {
             let result = a.asData(access: .noCopy)
-            XCTAssertEqual(result.shape, [4, 4])
-            XCTAssertEqual(result.strides, [4, 1])
-            XCTAssertEqual(result.dType, .int32)
+            #expect(result.shape == [4, 4])
+            #expect(result.strides == [4, 1])
+            #expect(result.dType == .int32)
         }
         do {
             let result = a.asData(access: .noCopyIfContiguous)
-            XCTAssertEqual(result.shape, [4, 4])
-            XCTAssertEqual(result.strides, [4, 1])
-            XCTAssertEqual(result.dType, .int32)
+            #expect(result.shape == [4, 4])
+            #expect(result.strides == [4, 1])
+            #expect(result.dType == .int32)
         }
     }
 
-    func testAsDataRoundTrip() {
+    @Test func asDataRoundTrip() {
         let a = MLXArray(0 ..< 16, [4, 4])
         let arrayData = a.asData(access: .copy)
         let result = MLXArray(arrayData.data, arrayData.shape, dtype: arrayData.dType)
-        assertEqual(a, result)
+        expectEqual(a, result)
     }
 
-    func testAsDataNonContiguous() {
+    @Test func asDataNonContiguous() {
         // buffer with holes (last dimension has stride of 2 and
         // thus larger storage than it physically needs)
         let a = MLXArray(0 ..< 16, [4, 4])
         let s = a[0..., .stride(by: 2)]
 
         let result = s.asData()
-        XCTAssertEqual(result.shape, [4, 2])
-        XCTAssertEqual(result.strides, [2, 1])
-        XCTAssertEqual(result.dType, .int32)
+        #expect(result.shape == [4, 2])
+        #expect(result.strides == [2, 1])
+        #expect(result.dType == .int32)
     }
 
-    func testAsDataNonContiguousNoCopy() {
+    @Test func asDataNonContiguousNoCopy() {
         // buffer with holes (last dimension has stride of 2 and
         // thus larger storage than it physically needs)
         let a = MLXArray(0 ..< 16, [4, 4])
@@ -196,24 +194,24 @@ class MLXArrayTests: XCTestCase {
         do {
             // the strides will match the strided array
             let result = s.asData(access: .noCopy)
-            XCTAssertEqual(result.shape, [4, 2])
-            XCTAssertEqual(result.strides, [4, 2])
-            XCTAssertEqual(result.dType, .int32)
+            #expect(result.shape == [4, 2])
+            #expect(result.strides == [4, 2])
+            #expect(result.dType == .int32)
         }
         do {
             // it isn't contiguous so we will get a contiguous copy
             let result = s.asData(access: .noCopyIfContiguous)
-            XCTAssertEqual(result.shape, [4, 2])
-            XCTAssertEqual(result.strides, [2, 1])
-            XCTAssertEqual(result.dType, .int32)
+            #expect(result.shape == [4, 2])
+            #expect(result.strides == [2, 1])
+            #expect(result.dType == .int32)
         }
     }
 
-    func testUpdateInternalWithItself() {
+    @Test func updateInternalWithItself() {
         // `Module.update(parameters:)` can hand an array its own value back
         let a = MLXArray([1, 2, 3] as [Float])
         a._updateInternal(a)
-        XCTAssertEqual(a.asArray(Float.self), [1, 2, 3])
+        #expect(a.asArray(Float.self) == [1, 2, 3])
     }
 
 }
diff --git a/Tests/MLXTests/MLXFastKernelTests.swift b/Tests/MLXTests/MLXFastKernelTests.swift
index 82f1dca8..4f637e07 100644
--- a/Tests/MLXTests/MLXFastKernelTests.swift
+++ b/Tests/MLXTests/MLXFastKernelTests.swift
@@ -3,13 +3,13 @@
 import Foundation
 import MLX
 import MLXNN
-import XCTest
+import Testing
 
-class MLXFastKernelTests: XCTestCase {
+@Suite
+struct MLXFastKernelTests {
 
-    func testCustomKernelBasic() {
+    @Test(.randomSeed(7)) func customKernelBasic() {
         // based on def test_custom_kernel_basic
-        MLXRandom.seed(7)
         let a = normal([2, 2])
         let kernel = MLXFast.metalKernel(
             name: "basic",
@@ -27,12 +27,11 @@ class MLXFastKernelTests: XCTestCase {
             outputShapes: [[2, 2]],
             outputDTypes: [.float32])
 
-        XCTAssertTrue(allClose(out[0], a).all().item())
+        expectEqual(out[0], a)
     }
 
-    func testCustomKernelArgs() {
+    @Test(.randomSeed(7)) func customKernelArgs() {
         // based on def test_custom_kernel_args
-        MLXRandom.seed(7)
         let a = normal([3, 6])
         let c = normal([2, 2]).asType(.bfloat16)
 
@@ -68,11 +67,11 @@ class MLXFastKernelTests: XCTestCase {
             outputShapes: [[2, 2], [3, 2]],
             outputDTypes: [.float32, .int32])
 
-        XCTAssertTrue(allClose(out[0], full([2, 2], values: 14.0484)).all().item())
-        XCTAssertTrue(allClose(out[1], full([3, 2], values: -2)).all().item())
+        expectEqual(out[0], full([2, 2], values: 14.0484))
+        expectEqual(out[1], full([3, 2], values: -2))
     }
 
-    func testFastSDPA() {
+    @Test func fastSDPA() {
         // https://github.com/ml-explore/mlx-swift/issues/172
         // this will just make sure the MLXFast.scaled_dot_product_attention is
         // callable in the various cases, based on
@@ -98,8 +97,7 @@ class MLXFastKernelTests: XCTestCase {
         }
     }
 
-    func testFastSDPAOutput() {
-        MLXRandom.seed(0)
+    @Test(.randomSeed(0)) func fastSDPAOutput() {
         let queries = MLXRandom.uniform(0.0 ..< 1.0, [1, 32, 1, 80])
         let keys = MLXRandom.uniform(0.0 ..< 1.0, [1, 32, 9, 80])
         let values = MLXRandom.uniform(0.0 ..< 1.0, [1, 32, 9, 80])
@@ -107,17 +105,16 @@ class MLXFastKernelTests: XCTestCase {
             queries: queries, keys: keys, values: values, scale: 0.1118, mask: .none)
         print(result.shape)
         print(result.sum().item(Float.self))
-        XCTAssertEqual(result.shape, [1, 32, 1, 80])
-        XCTAssertEqual(result.sum().item(Float.self), 1281.9253, accuracy: 0.01)
+        #expect(result.shape == [1, 32, 1, 80])
+        expectEqual(result.sum().item(Float.self), 1281.9253, accuracy: 0.01)
     }
 
-    func testRoPEOutput() {
+    @Test(.randomSeed(0)) func ropeOutput() {
         // https://github.com/ml-explore/mlx-swift/issues/315
-        MLXRandom.seed(0)
         let rope = RoPE(dimensions: 32, traditional: false, base: 10_000, scale: 1)
         let queries = MLXRandom.uniform(0.0 ..< 1.0, [1, 32, 1, 80])
         print(queries.shape)
         let result = rope(queries, offset: 8)
-        XCTAssertEqual(result.sum().item(Float.self), 1079.7894, accuracy: 0.01)
+        expectEqual(result.sum().item(Float.self), 1079.7894, accuracy: 0.01)
     }
 }
diff --git a/Tests/MLXTests/MLXRandomTests.swift b/Tests/MLXTests/MLXRandomTests.swift
index 43fa93c1..26fabd87 100644
--- a/Tests/MLXTests/MLXRandomTests.swift
+++ b/Tests/MLXTests/MLXRandomTests.swift
@@ -2,119 +2,117 @@
 
 import Foundation
 import MLX
-import XCTest
+import Testing
 
-class MLXRandomTests: XCTestCase {
+@Suite
+struct MLXRandomTests {
 
-    override class func setUp() {
-    }
-
-    func testSplit() {
+    @Test func splitKey() {
         let key = MLXRandom.key(0)
         let keys = MLXRandom.split(key: key, into: 4)
 
-        XCTAssertEqual(keys.count, 4)
+        #expect(keys.count == 4)
     }
 
-    func testUniformSingle() {
+    @Test func uniformSingle() {
         let key = MLXRandom.key(0)
         let value = MLXRandom.uniform(0 ..< 10, key: key).item(Float.self)
-        XCTAssertEqual(value, 4.18, accuracy: 0.01)
+        expectEqual(value, 4.18, accuracy: 0.01)
     }
 
-    func testUniformMultiple() {
+    @Test func uniformMultiple() {
         let key = MLXRandom.key(0)
 
         // specify shape to broadcast
         let value = MLXRandom.uniform(0.0 ..< 10, [3], key: key)
         let expected = MLXArray(converting: [9.65, 3.14, 6.33])
-        assertEqual(value, expected, atol: 0.01)
+        expectEqual(value, expected, atol: 0.01)
     }
 
-    func testUniformMultipleArray() {
+    @Test func uniformMultipleArray() {
         let key = MLXRandom.key(0)
 
         // give an array of bounds
         let value = MLXRandom.uniform(low: [0, 10], high: [10, 100], key: key)
 
         let expected = MLXArray(converting: [2.16, 82.37])
-        assertEqual(value, expected, atol: 0.01)
+        expectEqual(value, expected, atol: 0.01)
     }
 
-    func testNormal() {
+    @Test func normalSingle() {
         let key = MLXRandom.key(0)
 
         let value = MLXRandom.normal(key: key).item(Float.self)
-        XCTAssertEqual(value, -0.20, accuracy: 0.01)
+        expectEqual(value, -0.20, accuracy: 0.01)
     }
 
-    func testRandIntSingle() {
+    @Test func randIntSingle() {
         let key = MLXRandom.key(0)
         let value = MLXRandom.randInt(0 ..< 100, key: key).item(Int.self)
-        XCTAssertEqual(value, 41)
+        #expect(value == 41)
     }
 
-    func testRandIntMultiple() {
+    @Test func randIntMultiple() {
         let key = MLXRandom.key(0)
 
         let value = MLXRandom.randInt(low: [0, 10], high: [10, 100], key: key)
 
         let expected: [Int32] = [2, 82]
-        assertEqual(value, MLXArray(expected))
+        expectEqual(value, MLXArray(expected))
     }
 
-    func testRandIntMultipleType() {
+    @Test func randIntMultipleType() {
         let key = MLXRandom.key(0)
 
         let value = MLXRandom.randInt(low: [0, 10], high: [10, 100], type: Int8.self, key: key)
 
         let expected: [Int8] = [2, 82]
-        assertEqual(value, MLXArray(expected))
+        expectEqual(value, MLXArray(expected))
     }
 
-    func testBernoulliSingle() {
+    @Test func bernoulliSingle() {
         let key = MLXRandom.key(0)
         let value = MLXRandom.bernoulli(key: key).item(Bool.self)
-        XCTAssertEqual(value, true)
+        #expect(value == true)
     }
 
-    func testBernoulliMultiple() {
+    @Test func bernoulliMultiple() {
         let key = MLXRandom.key(0)
         let array = MLXRandom.bernoulli([4], key: key)
         let expected = MLXArray([false, true, false, true])
-        assertEqual(array, expected)
+        expectEqual(array, expected)
     }
 
-    func testBernoulliP() {
+    @Test func bernoulliP() {
         let key = MLXRandom.key(0)
         let array = MLXRandom.bernoulli(0.8, [4], key: key)
         let expected = MLXArray([false, true, true, true])
-        assertEqual(array, expected)
+        expectEqual(array, expected)
     }
 
-    func testBernoulliPArray() {
+    @Test func bernoulliPArray() {
         let key = MLXRandom.key(0)
         let array = MLXRandom.bernoulli(MLXArray(converting: [0.1, 0.5, 0.8]), key: key)
         let expected = MLXArray([false, true, true])
-        assertEqual(array, expected)
+        expectEqual(array, expected)
     }
 
-    func testTruncatedNormalSingle() {
+    @Test func truncatedNormalSingle() {
         let key = MLXRandom.key(0)
         let value = MLXRandom.truncatedNormal(0 ..< 10, key: key).item(Float.self)
-        XCTAssertEqual(value, 0.55, accuracy: 0.01)
+        expectEqual(value, 0.55, accuracy: 0.01)
     }
 
-    func testTruncatedNormalMultiple() {
+    @Test func truncatedNormalMultiple() {
         let key = MLXRandom.key(0)
 
         // specify shape to broadcast
         let value = MLXRandom.truncatedNormal(0 ..< 0.5, [3], key: key)
         let expected = MLXArray(converting: [0.48, 0.15, 0.30])
-        assertEqual(value, expected, atol: 0.01)
+        expectEqual(value, expected, atol: 0.01)
     }
 
-    func testTruncatedNormalMultipleArray() {
+    @Test func truncatedNormalMultipleArray() {
         let key = MLXRandom.key(0)
 
         // give an array of bounds
@@ -122,54 +120,60 @@ class MLXRandomTests: XCTestCase {
             low: MLXArray(converting: [0, 0.5]), high: MLXArray(converting: [0.5, 1]), key: key)
 
         let expected = MLXArray(converting: [0.10, 0.88])
-        assertEqual(value, expected, atol: 0.01)
+        expectEqual(value, expected, atol: 0.01)
     }
 
-    func testGumbel() {
+    @Test func gumbelSingle() {
         let key = MLXRandom.key(0)
 
         let value = MLXRandom.gumbel(key: key).item(Float.self)
-        XCTAssertEqual(value, 0.13, accuracy: 0.01)
+        expectEqual(value, 0.13, accuracy: 0.01)
     }
 
-    func testLogits() {
+    @Test func categoricalLogits() {
         let key = MLXRandom.key(0)
 
         let logits = MLXArray.zeros([5, 20])
         let result = MLXRandom.categorical(logits, key: key)
 
-        XCTAssertEqual(result.shape, [5])
-        XCTAssertEqual(result.dtype, .uint32)
+        #expect(result.shape == [5])
+        #expect(result.dtype == .uint32)
 
         let expected = MLXArray([1, 1, 17, 17, 17])
-        assertEqual(result, expected)
+        expectEqual(result, expected)
     }
 
-    func testLogitsCount() {
+    @Test func categoricalLogitsCount() {
         let key = MLXRandom.key(0)
 
         let logits = MLXArray.zeros([5, 20])
         let result = MLXRandom.categorical(logits, count: 2, key: key)
 
-        XCTAssertEqual(result.shape, [5, 2])
-        XCTAssertEqual(result.dtype, .uint32)
+        #expect(result.shape == [5, 2])
+        #expect(result.dtype == .uint32)
 
         let expected = MLXArray([16, 3, 14, 10, 17, 7, 6, 8, 12, 8], [5, 2])
-        assertEqual(result, expected)
+        expectEqual(result, expected)
     }
 
-    func testRandomStateOrKeySame() {
+    @Test func randomStateOrKeySame() {
         // these should all produce the same value since they
         // all resolve to the same key
+        //
+        // note: this does not use the global state (MLXRandom.seed()) -- it is
+        // shared by every test in the process and other tests running in parallel
+        // may draw from it.  MLXRandom.seed(n) is MLXRandom.globalState.seed(n),
+        // which is covered by v3 below.
 
         let key = MLXRandom.key(0)
         let (_, k1) = split(key: key)
 
         let state = MLXRandom.RandomState(seed: 0)
-        MLXRandom.seed(0)
 
-        // global state
-        let v0 = uniform(0 ..< 1, [5])
+        // task local RandomState (the fallback when no key is given)
+        let v0 = withRandomState(MLXRandom.RandomState(seed: 0)) {
+            uniform(0 ..< 1, [5])
+        }
 
         // explicit key
         let v1 = uniform(0 ..< 1, [5], key: k1)
@@ -177,23 +181,32 @@ class MLXRandomTests: XCTestCase {
         // local RandomState
         let v2 = uniform(0 ..< 1, [5], key: state)
 
-        assertEqual(v0, v1)
-        assertEqual(v1, v2)
+        // a RandomState that is re-seeded, the same as MLXRandom.seed(0) does
+        // to the global state
+        let reseeded = MLXRandom.RandomState(seed: 99)
+        _ = uniform(0 ..< 1, [5], key: reseeded)
+        reseeded.seed(0)
+        let v3 = uniform(0 ..< 1, [5], key: reseeded)
+
+        expectEqual(v0, v1)
+        expectEqual(v1, v2)
+        expectEqual(v2, v3)
     }
 
-    func testRandomStateOrKeyDifferent() {
+    @Test func randomStateOrKeyDifferent() {
         // these should all produce different values as they
         // use different keys -- note this is otherwise identical
-        // to testRandomStateOrKeySame
+        // to randomStateOrKeySame
 
         let key = MLXRandom.key(7)
         let (_, k1) = split(key: key)
 
         let state = MLXRandom.RandomState(seed: 11)
-        MLXRandom.seed(31)
 
-        // global state
-        let v0 = uniform(0 ..< 1, [5])
+        // task local RandomState
+        let v0 = withRandomState(MLXRandom.RandomState(seed: 31)) {
+            uniform(0 ..< 1, [5])
+        }
 
         // explicit key
         let v1 = uniform(0 ..< 1, [5], key: k1)
@@ -201,12 +214,12 @@ class MLXRandomTests: XCTestCase {
         // local RandomState
         let v2 = uniform(0 ..< 1, [5], key: state)
 
-        assertNotEqual(v0, v1)
-        assertNotEqual(v1, v2)
-        assertNotEqual(v0, v2)
+        expectNotEqual(v0, v1)
+        expectNotEqual(v1, v2)
+        expectNotEqual(v0, v2)
     }
 
-    func testRandomThreadsSame() async {
+    @Test func randomThreadsSame() async {
         // several threads using task local random state with a constant
         // seed will produce the same value
         await withTaskGroup(of: Float.self) { group in
@@ -229,11 +242,11 @@ class MLXRandomTests: XCTestCase {
             }
 
             let unique = Set(result)
-            XCTAssertEqual(unique.count, 1, "Different values: \(result)")
+            #expect(unique.count == 1, "Different values: \(result)")
         }
     }
 
-    func testRandomThreadsDifferent() async {
+    @Test func randomThreadsDifferent() async {
 
         // several threads using task local random state with different
         // seeds will produce different values
@@ -259,7 +272,7 @@ class MLXRandomTests: XCTestCase {
             }
 
             let unique = Set(result)
-            XCTAssertEqual(unique.count, 10, "Same values: \(result)")
+            #expect(unique.count == 10, "Same values: \(result)")
         }
     }
 
diff --git a/Tests/MLXTests/MaterializedHazardTests.swift b/Tests/MLXTests/MaterializedHazardTests.swift
index 644bb9c6..89b6d3c4 100644
--- a/Tests/MLXTests/MaterializedHazardTests.swift
+++ b/Tests/MLXTests/MaterializedHazardTests.swift
@@ -129,7 +129,7 @@ private func compileSealedModuleWhileShared() {
                 eval(compiled(x))
             } else {
                 let r = mm(x)
-                #expect(r.allClose(expected).item(Bool.self))
+                expectEqual(r, expected)
             }
         }
     }
@@ -241,7 +241,7 @@ struct MaterializedHazardTests {
         let x = MLXRandom.normal([4, 16])
         let expected = mm(x)
         let result = compiled(x)
-        #expect(result.allClose(expected).item(Bool.self))
+        expectEqual(result, expected)
 
         // make sure the trace actually ran so the checks below mean something
         #expect(traced)
@@ -381,7 +381,7 @@ private final class ReleasedCompileState: @unchecked Sendable {
             var output = mlx_vector_array_new()
             defer { mlx_vector_array_free(output) }
             mlx_closure_apply(&output, replacement, input)
-            return mlx_vector_array_values(output)[0]
+            return mlx_vector_array_values(output, count: 1)[0]
         }
         return result.item(Float.self)
     }
@@ -396,7 +396,8 @@ struct CompileReboundMaterializedStateTests {
         let released = DispatchSemaphore(value: 0)
         let finished = DispatchSemaphore(value: 0)
         let worker = Thread {
-            state.before = state.compiled!.call([MLXArray(Float(2))])[0].item(Float.self)
+            state.before = state.compiled!.call([MLXArray(Float(2))], expected: 1)[0]
+                .item(Float.self)
             primed.signal()
             if released.wait(timeout: .now() + 30) == .success {
                 state.after = state.replacementValue()
@@ -459,10 +460,10 @@ struct CompileReboundMaterializedStateTests {
         let holder = ModuleHolder(makeLinear(1, sealed: false))
         let c = compile(inputs: [holder], outputs: [holder]) { (x: MLXArray) in holder.model(x) }
 
-        #expect(c(input).allClose(holder.model(input)).item(Bool.self))
+        expectEqual(c(input), holder.model(input))
 
         holder.model = makeLinear(2, sealed: false)
-        #expect(c(input).allClose(holder.model(input)).item(Bool.self))
+        expectEqual(c(input), holder.model(input))
     }
 
     /// Control: a plain array rebound between calls.
@@ -471,10 +472,10 @@ struct CompileReboundMaterializedStateTests {
         let holder = ArrayHolder(MLXArray(Float(1)))
         let c = compile(inputs: [holder], outputs: [holder]) { (x: MLXArray) in x + holder.value }
 
-        #expect(c(input).allClose(input + 1).item(Bool.self))
+        expectEqual(c(input), input + 1)
 
         holder.value = MLXArray(Float(2))
-        #expect(c(input).allClose(input + 2).item(Bool.self))
+        expectEqual(c(input), input + 2)
     }
 
     /// Sealed A -> sealed A2 (same shapes): both have zero threaded state
@@ -485,11 +486,11 @@ struct CompileReboundMaterializedStateTests {
         let holder = ModuleHolder(makeLinear(1, sealed: true))
         let c = compile(inputs: [holder], outputs: [holder]) { (x: MLXArray) in holder.model(x) }
 
-        #expect(c(input).allClose(holder.model(input)).item(Bool.self))
+        expectEqual(c(input), holder.model(input))
 
         holder.model = makeLinear(2, sealed: true)
-        #expect(
-            c(input).allClose(holder.model(input)).item(Bool.self),
+        expectEqual(
+            c(input), holder.model(input),
             "compiled function used the weights of the module it was traced with")
     }
 
@@ -501,14 +502,14 @@ struct CompileReboundMaterializedStateTests {
         let holder = ModuleHolder(makeLinear(1, sealed: true))
         let c = compile(inputs: [holder], outputs: [holder]) { (x: MLXArray) in holder.model(x) }
 
-        #expect(c(input).allClose(holder.model(input)).item(Bool.self))
+        expectEqual(c(input), holder.model(input))
 
         holder.model = makeLinear(2, sealed: false)
-        #expect(c(input).allClose(holder.model(input)).item(Bool.self))
+        expectEqual(c(input), holder.model(input))
 
         holder.model = makeLinear(3, sealed: true)
-        #expect(
-            c(input).allClose(holder.model(input)).item(Bool.self),
+        expectEqual(
+            c(input), holder.model(input),
             "compiled function reused the trace of an earlier sealed module")
     }
 
@@ -520,11 +521,11 @@ struct CompileReboundMaterializedStateTests {
         let holder = ModuleHolder(partlySealed(1))
         let c = compile(inputs: [holder], outputs: [holder]) { (x: MLXArray) in holder.model(x) }
 
-        #expect(c(input).allClose(holder.model(input)).item(Bool.self))
+        expectEqual(c(input), holder.model(input))
 
         holder.model = partlySealed(2)
-        #expect(
-            c(input).allClose(holder.model(input)).item(Bool.self),
+        expectEqual(
+            c(input), holder.model(input),
             "compiled function mixed baked weights from the traced module")
     }
 
@@ -536,14 +537,14 @@ struct CompileReboundMaterializedStateTests {
         let parent = Parent(makeLinear(1, sealed: true))
         let c = compile(inputs: [parent], outputs: [parent]) { (x: MLXArray) in parent(x) }
 
-        #expect(c(input).allClose(parent(input)).item(Bool.self))
+        expectEqual(c(input), parent(input))
 
         var u = NestedDictionary<String, Module>()
         u["child"] = .value(makeLinear(2, sealed: true))
         parent.update(modules: u)
 
-        #expect(
-            c(input).allClose(parent(input)).item(Bool.self),
+        expectEqual(
+            c(input), parent(input),
             "compiled function used the weights of the child it was traced with")
     }
 
@@ -556,11 +557,11 @@ struct CompileReboundMaterializedStateTests {
         let holder = ArrayHolder(MLXArray(Float(1)).materialized())
         let c = compile(inputs: [holder], outputs: [holder]) { (x: MLXArray) in x + holder.value }
 
-        #expect(c(input).allClose(input + 1).item(Bool.self))
+        expectEqual(c(input), input + 1)
 
         holder.value = MLXArray(Float(2)).materialized()
-        #expect(
-            c(input).allClose(input + 2).item(Bool.self),
+        expectEqual(
+            c(input), input + 2,
             "compiled function used the MaterializedArray it was traced with")
     }
 
@@ -583,7 +584,7 @@ struct CompileReboundMaterializedStateTests {
 
         holder.model = makeLinear(2, sealed: true)
         for _ in 0 ..< 3 {
-            #expect(c(input).allClose(holder.model(input)).item(Bool.self))
+            expectEqual(c(input), holder.model(input))
         }
         #expect(traces == 2)
     }
@@ -641,15 +642,12 @@ struct SendableTransformsMaterializedTests {
         let input = compileInput().materialized()
         let expected = model(input).materialized()
 
-        await withTaskGroup(of: Bool.self) { group in
+        await withTaskGroup(of: Void.self) { group in
             for _ in 0 ..< 8 {
                 group.addTask {
-                    compiled(input).allClose(expected).item(Bool.self)
+                    expectEqual(compiled(input), expected)
                 }
             }
-            for await ok in group {
-                #expect(ok)
-            }
         }
     }
 
@@ -661,7 +659,7 @@ struct SendableTransformsMaterializedTests {
         }
         for offset in 0 ..< 3 {
             let input = compileInput() + Float(offset)
-            #expect(compiled(input).allClose(input * 3).item(Bool.self))
+            expectEqual(compiled(input), input * 3)
             #expect(isEvaluated(scale))
             #expect(!hasGraph(scale))
             #expect(scale.item(Float.self) == 3)
@@ -674,21 +672,18 @@ struct SendableTransformsMaterializedTests {
             x * scale
         }
 
-        await withTaskGroup(of: Bool.self) { group in
+        await withTaskGroup(of: Void.self) { group in
             for worker in 0 ..< 8 {
                 group.addTask {
                     for iteration in 0 ..< 25 {
                         let input = compileInput() + Float(worker * 25 + iteration)
-                        if !compiled(input).allClose(input * 3).item(Bool.self) {
-                            return false
+                        // stop at the first failure rather than recording 25 of them
+                        if !expectEqual(compiled(input), input * 3) {
+                            return
                         }
                     }
-                    return true
                 }
             }
-            for await ok in group {
-                #expect(ok)
-            }
         }
         #expect(isEvaluated(scale))
         #expect(!hasGraph(scale))
diff --git a/Tests/MLXTests/MaterializedTests.swift b/Tests/MLXTests/MaterializedTests.swift
index cdc55f18..4f0da61d 100644
--- a/Tests/MLXTests/MaterializedTests.swift
+++ b/Tests/MLXTests/MaterializedTests.swift
@@ -34,7 +34,7 @@ struct MaterializedTests {
         let t = Task {
             let s = sum(c(i1, i2))
             let s2 = sum(f(i1, i2))
-            #expect(s.allClose(s2).item(Bool.self))
+            expectEqual(s, s2)
             print(s, s2)
         }
         _ = await t.result
@@ -57,7 +57,7 @@ struct MaterializedTests {
         let x = MLXRandom.normal([4, 4])
         let r1 = compiled(x)
         let r2 = x + state
-        #expect(r1.allClose(r2).item(Bool.self))
+        expectEqual(r1, r2)
 
         // same for a `Module` whose parameters were materialized (e.g. via
         // `MaterializedModule`) -- this is the documented
@@ -75,7 +75,7 @@ struct MaterializedTests {
         let input = MLXRandom.normal([4, 4])
         let modelResult = compiledModel(input)
         let expected = l(input)
-        #expect(modelResult.allClose(expected).item(Bool.self))
+        expectEqual(modelResult, expected)
     }
 
     @Test
diff --git a/Tests/MLXTests/MemoryTests.swift b/Tests/MLXTests/MemoryTests.swift
index dc16ca9b..1f3163d6 100644
--- a/Tests/MLXTests/MemoryTests.swift
+++ b/Tests/MLXTests/MemoryTests.swift
@@ -2,30 +2,45 @@
 
 import Foundation
 import MLX
-import XCTest
+import Testing
 
-class MemoryTests: XCTestCase {
+@Suite
+struct MemoryTests {
 
-    func testWiredMemory() {
-        Memory.withWiredLimit(1024 * 1024 * 256) {
+    @Test func wiredMemoryDeprecated() {
+        // this is a NOP wired-memory-wise but should still work
+        let r = Memory.withWiredLimit(1024 * 1024 * 256) {
             let x = MLXArray(10)
-            print(x * x)
+            return x * 3
         }
+
+        #expect(r.item(Int.self) == 30)
+    }
+
+    @Test func wiredMemory() async {
+        let policy = WiredSumPolicy()
+        let weights = policy.ticket(size: 1024 * 1024)
+        let r = await weights.withWiredLimit {
+            let x = MLXArray(10)
+            return x * 3
+        }
+
+        #expect(r.item(Int.self) == 30)
     }
 
-    func testDeviceInfo() {
+    @Test func deviceInfo() {
         let info = GPU.deviceInfo()
-        XCTAssertGreaterThan(info.memorySize, 0)
+        #expect(info.memorySize > 0)
 
         guard Device.defaultDevice().deviceType == .gpu else { return }
-        XCTAssertNotEqual(info.architecture, "Unknown")
-        XCTAssertGreaterThan(info.maxBufferSize, 0)
-        XCTAssertGreaterThan(info.maxRecommendedWorkingSetSize, 0)
+        #expect(info.architecture != "Unknown")
+        #expect(info.maxBufferSize > 0)
+        #expect(info.maxRecommendedWorkingSetSize > 0)
 
         let again = GPU.deviceInfo()
-        XCTAssertEqual(again.architecture, info.architecture)
-        XCTAssertEqual(again.maxBufferSize, info.maxBufferSize)
-        XCTAssertEqual(again.maxRecommendedWorkingSetSize, info.maxRecommendedWorkingSetSize)
+        #expect(again.architecture == info.architecture)
+        #expect(again.maxBufferSize == info.maxBufferSize)
+        #expect(again.maxRecommendedWorkingSetSize == info.maxRecommendedWorkingSetSize)
     }
 
     // `Memory.bufferSize` reports the allocator length of each buffer, and the
@@ -34,27 +49,27 @@ class MemoryTests: XCTestCase {
     // tests running in parallel).  These tests only check the properties that hold
     // regardless of the cache.
 
-    func testBufferSizeAtLeastLogicalSize() {
+    @Test func bufferSizeAtLeastLogicalSize() {
         let a = MLXRandom.normal([64, 32])
         let b = MLXArray.zeros([10, 10], type: UInt8.self)
         eval(a, b)
 
-        XCTAssertGreaterThanOrEqual(Memory.bufferSize([a]), a.nbytes)
-        XCTAssertGreaterThanOrEqual(Memory.bufferSize([b]), b.nbytes)
+        #expect(Memory.bufferSize([a]) >= a.nbytes)
+        #expect(Memory.bufferSize([b]) >= b.nbytes)
     }
 
-    func testBufferSizeEmpty() {
-        XCTAssertEqual(Memory.bufferSize([MLXArray]()), 0)
+    @Test func bufferSizeEmpty() {
+        #expect(Memory.bufferSize([MLXArray]()) == 0)
     }
 
-    func testBufferSizeCountsDuplicatesOnce() {
+    @Test func bufferSizeCountsDuplicatesOnce() {
         let a = MLXRandom.normal([64, 32])
         eval(a)
 
-        XCTAssertEqual(Memory.bufferSize([a, a, a]), Memory.bufferSize([a]))
+        #expect(Memory.bufferSize([a, a, a]) == Memory.bufferSize([a]))
     }
 
-    func testBufferSizeCountsViewsOnce() {
+    @Test func bufferSizeCountsViewsOnce() {
         let a = MLXRandom.normal([64, 32])
         let slice = a[0 ..< 8]
         let transposed = a.T
@@ -62,19 +77,19 @@ class MemoryTests: XCTestCase {
 
         // views share the buffer of `a`
         let size = Memory.bufferSize([a])
-        XCTAssertEqual(Memory.bufferSize([a, slice]), size)
-        XCTAssertEqual(Memory.bufferSize([a, transposed]), size)
-        XCTAssertEqual(Memory.bufferSize([slice, transposed]), size)
+        #expect(Memory.bufferSize([a, slice]) == size)
+        #expect(Memory.bufferSize([a, transposed]) == size)
+        #expect(Memory.bufferSize([slice, transposed]) == size)
     }
 
-    func testBufferSizeSumsDistinctBuffers() {
+    @Test func bufferSizeSumsDistinctBuffers() {
         let a = MLXRandom.normal([64, 32])
         let b = MLXRandom.normal([5])
         let c = MLXRandom.randInt(low: 0, high: 4, [100], type: UInt32.self)
         eval(a, b, c)
 
-        XCTAssertEqual(
-            Memory.bufferSize([a, b, c]),
-            Memory.bufferSize([a]) + Memory.bufferSize([b]) + Memory.bufferSize([c]))
+        #expect(
+            Memory.bufferSize([a, b, c]) == Memory.bufferSize([a]) + Memory.bufferSize([b])
+                + Memory.bufferSize([c]))
     }
 }
diff --git a/Tests/MLXTests/ModuleTests.swift b/Tests/MLXTests/ModuleTests.swift
index c314addf..aa104033 100644
--- a/Tests/MLXTests/ModuleTests.swift
+++ b/Tests/MLXTests/ModuleTests.swift
@@ -2,14 +2,12 @@
 
 import Foundation
 import MLX
-import XCTest
+import Testing
 
 @testable import MLXNN
 
-class ModuleTests: XCTestCase {
-
-    override class func setUp() {
-    }
+@Suite
+struct ModuleTests {
 
     func newTestModule() -> Module {
         class ChildTestModule: Module {
@@ -51,7 +49,7 @@ class ModuleTests: XCTestCase {
         return TestModule()
     }
 
-    func testModuleDescription() {
+    @Test func moduleDescription() {
         let m = newTestModule()
 
         let expected =
@@ -66,10 +64,10 @@ class ModuleTests: XCTestCase {
             """
 
         let d = m.description
-        XCTAssertEqual(d, expected)
+        #expect(d == expected)
     }
 
-    func testNested() {
+    @Test func nested() {
         let m = newTestModule()
         let v = m.filterMap(filter: Module.filterAll)
 
@@ -123,18 +121,18 @@ class ModuleTests: XCTestCase {
             ]
             """
 
-        XCTAssertEqual(v.count, 6)
+        #expect(v.count == 6)
         let d = v.description
-        XCTAssertEqual(d, expected)
+        #expect(d == expected)
     }
 
-    func testModuleItems() {
+    @Test func moduleItems() {
         let m = newTestModule()
 
         // all the ivars from TestModule
         let t = m.items()
 
-        XCTAssertEqual(t.count, 6)
+        #expect(t.count == 6)
 
         for (k, v) in t {
             switch (k, v) {
@@ -145,16 +143,16 @@ class ModuleTests: XCTestCase {
             case ("alternateNameArray", .array): break
             case ("alternateNameDict", .dictionary): break
             default:
-                XCTFail("Unexpected: \(k) = \(v)))")
+                Issue.record("Unexpected: \(k) = \(v)))")
             }
         }
     }
 
-    func testModuleParameters() {
+    @Test func moduleParameters() {
         let m = newTestModule()
 
         let p = m.parameters()
-        XCTAssertEqual(p.count, 5)
+        #expect(p.count == 5)
 
         for (k, v) in p {
             switch (k, v) {
@@ -165,12 +163,12 @@ class ModuleTests: XCTestCase {
             case ("d", .dictionary): break
             case ("alternateChild", .dictionary): break
             default:
-                XCTFail("Unexpected: \(k) = \(v)))")
+                Issue.record("Unexpected: \(k) = \(v)))")
             }
         }
     }
 
-    func testTupleParameters() {
+    @Test func tupleParameters() {
         // make sure tuples of MLXArray work
         class Test: Module {
             let tuple: (MLXArray, MLXArray)
@@ -185,18 +183,18 @@ class ModuleTests: XCTestCase {
         let t = Test()
         let p = t.parameters()
 
-        XCTAssertEqual(p.count, 2)
+        #expect(p.count == 2)
         switch p["array"] {
         case .array(let items):
-            XCTAssertEqual(items.count, 2)
+            #expect(items.count == 2)
         default:
-            XCTFail("should be a .array with 2 MLXArray: \(String(describing: p["item"]))")
+            Issue.record("should be a .array with 2 MLXArray: \(String(describing: p["item"]))")
         }
         switch p["tuple"] {
         case .array(let items):
-            XCTAssertEqual(items.count, 2)
+            #expect(items.count == 2)
         default:
-            XCTFail("should be a .array with 2 MLXArray: \(String(describing: p["item"]))")
+            Issue.record("should be a .array with 2 MLXArray: \(String(describing: p["item"]))")
         }
 
         // update the second item on the array
@@ -206,10 +204,10 @@ class ModuleTests: XCTestCase {
 
         // make sure it was written
         if let v1 = t.parameters()["array"]?.flattenedValues() {
-            XCTAssertEqual(v1.count, 2)
-            XCTAssertEqual(v1[1].item(Int.self), 7)
+            #expect(v1.count == 2)
+            #expect(v1[1].item(Int.self) == 7)
         } else {
-            XCTFail("unable to read array")
+            Issue.record("unable to read array")
         }
 
         // update the second item on the tuple
@@ -219,14 +217,14 @@ class ModuleTests: XCTestCase {
 
         // make sure it was written
         if let v2 = t.parameters()["tuple"]?.flattenedValues() {
-            XCTAssertEqual(v2.count, 2)
-            XCTAssertEqual(v2[1].item(Int.self), 11)
+            #expect(v2.count == 2)
+            #expect(v2[1].item(Int.self) == 11)
         } else {
-            XCTFail("unable to read tuple")
+            Issue.record("unable to read tuple")
         }
     }
 
-    func testTupleModules() {
+    @Test func tupleModules() {
         class Layer1: Module {
             let x: MLXArray
 
@@ -255,7 +253,7 @@ class ModuleTests: XCTestCase {
         let t = Test()
 
         // top level + 2 + 2
-        XCTAssertEqual(t.modules().count, 5)
+        #expect(t.modules().count == 5)
 
         // update the second item on the array
         var u1 = ModuleParameters()
@@ -267,10 +265,10 @@ class ModuleTests: XCTestCase {
 
         // make sure it was written
         if let v1 = t.parameters()["array"]?.flattenedValues() {
-            XCTAssertEqual(v1.count, 2)
-            XCTAssertEqual(v1[1].item(Int.self), 7)
+            #expect(v1.count == 2)
+            #expect(v1[1].item(Int.self) == 7)
         } else {
-            XCTFail("unable to read array")
+            Issue.record("unable to read array")
         }
 
         // update the second item on the tuple
@@ -283,17 +281,17 @@ class ModuleTests: XCTestCase {
 
         // make sure it was written
         if let v2 = t.parameters()["tuple"]?.flattenedValues() {
-            XCTAssertEqual(v2.count, 2)
-            XCTAssertEqual(v2[1].item(Int.self), 11)
+            #expect(v2.count == 2)
+            #expect(v2[1].item(Int.self) == 11)
         } else {
-            XCTFail("unable to read tuple")
+            Issue.record("unable to read tuple")
         }
 
         print(t.parameters())
         print(t.parameters().flattened())
     }
 
-    func testOptionInfos() {
+    @Test func optionInfos() {
         class Layer1: Module {
             @ParameterInfo var x: MLXArray?
         }
@@ -309,28 +307,27 @@ class ModuleTests: XCTestCase {
 
         let t = Test()
 
-        XCTAssertNil(t.b)
-        XCTAssertNotNil(t.a)
-        XCTAssertNil(t.a.x)
+        #expect(t.b == nil)
+        #expect(t.a.x == nil)
     }
 
-    func testChildren() {
+    @Test func children() {
         let m = newTestModule()
 
         let c = m.children()
-        XCTAssertEqual(c.count, 2)
+        #expect(c.count == 2)
 
         for (k, v) in c {
             switch (k, v) {
             case ("alternateChild", .value): break
             case ("d", .dictionary): break
             default:
-                XCTFail("Unexpected: \(k) = \(v)))")
+                Issue.record("Unexpected: \(k) = \(v)))")
             }
         }
     }
 
-    func testSequentialBuilder() {
+    @Test func sequentialBuilder() {
         let b = Bool.random()
         let s = Sequential {
             Tanh()
@@ -343,10 +340,10 @@ class ModuleTests: XCTestCase {
                 Linear(10, 20)
             }
         }
-        XCTAssertEqual(s.layers.count, 5)
+        #expect(s.layers.count == 5)
     }
 
-    func testInit() {
+    @Test func propertyWrapperInit() {
         // ensure that different types of vars and init work
 
         class Test: Module {
@@ -366,10 +363,10 @@ class ModuleTests: XCTestCase {
         let t = Test()
 
         // 2 + 2 * 2 (Linear)
-        XCTAssertEqual(t.parameters().flattenedValues().count, 6)
+        #expect(t.parameters().flattenedValues().count == 6)
 
         // self + 2 Linears
-        XCTAssertEqual(t.modules().count, 3)
+        #expect(t.modules().count == 3)
     }
 
     func newStructureModule() -> Module {
@@ -406,32 +403,32 @@ class ModuleTests: XCTestCase {
         return StructureModel()
     }
 
-    func testLeafModules() {
+    @Test func leafModules() {
         let m = newStructureModule()
         let modules = m.leafModules()
         let v = Dictionary(uniqueKeysWithValues: modules.flattened())
 
-        XCTAssertEqual(v.count, 6)
-        XCTAssertEqual(
-            Set(v.keys),
-            Set([
-                "child", "interior2.children.1", "leaf", "interior2.children.0", "interior1.child",
-                "nonWrappedChild",
-            ]))
+        #expect(v.count == 6)
+        #expect(
+            Set(v.keys)
+                == Set([
+                    "child", "interior2.children.1", "leaf", "interior2.children.0",
+                    "interior1.child",
+                    "nonWrappedChild",
+                ]))
     }
 
-    func testChildren2() {
+    @Test func children2() {
         let m = newStructureModule()
         // only the modules attached to the top level
         let modules = m.children()
         let v = Dictionary(uniqueKeysWithValues: modules.flattened())
 
-        XCTAssertEqual(v.count, 5)
-        XCTAssertEqual(
-            Set(v.keys), Set(["child", "interior2", "leaf", "interior1", "nonWrappedChild"]))
+        #expect(v.count == 5)
+        #expect(Set(v.keys) == Set(["child", "interior2", "leaf", "interior1", "nonWrappedChild"]))
     }
 
-    func testVisitModule() {
+    @Test func visitModule() {
         let m = newStructureModule()
         var collect = [String: String]()
 
@@ -439,12 +436,12 @@ class ModuleTests: XCTestCase {
             collect[key] = String(describing: type(of: m))
         }
 
-        XCTAssertEqual(collect.count, 9)
-        XCTAssertEqual(collect[""], "StructureModel")
-        XCTAssertEqual(collect["interior2.children.1"], "Linear")
+        #expect(collect.count == 9)
+        #expect(collect[""] == "StructureModel")
+        #expect(collect["interior2.children.1"] == "Linear")
     }
 
-    func testUpdateParametersCallsChildOverride() throws {
+    @Test func updateParametersCallsChildOverride() throws {
         // a child's update(parameters:verify:path:modulePath:) override runs
         // when the update starts at an ancestor
         class RecordingLinear: Linear {
@@ -472,13 +469,13 @@ class ModuleTests: XCTestCase {
                 "child.weight": weight, "list.0.weight": weight,
             ]))
 
-        XCTAssertEqual(m.child.updateCount, 1)
-        XCTAssertEqual(m.list[0].updateCount, 1)
-        assertEqual(m.child.weight, weight)
-        assertEqual(m.list[0].weight, weight)
+        #expect(m.child.updateCount == 1)
+        #expect(m.list[0].updateCount == 1)
+        expectEqual(m.child.weight, weight)
+        expectEqual(m.list[0].weight, weight)
     }
 
-    func testUpdateModuleSameType() throws {
+    @Test func updateModuleSameType() throws {
         // set a single module with same type
         let m = newStructureModule()
         var u = NestedDictionary<String, Module>()
@@ -488,13 +485,13 @@ class ModuleTests: XCTestCase {
         let modules = m.children()
         if let child = modules["child"], let child = child.unwrap() as? Linear {
             // make sure it was updated
-            XCTAssertEqual(child.shape.0, 4)
+            #expect(child.shape.0 == 4)
         } else {
-            XCTFail("child is nil or not Linear")
+            Issue.record("child is nil or not Linear")
         }
     }
 
-    func testUpdateModuleSubType() throws {
+    @Test func updateModuleSubType() throws {
         // set a single module with a subtype
         let m = newStructureModule()
         var u = NestedDictionary<String, Module>()
@@ -504,13 +501,13 @@ class ModuleTests: XCTestCase {
         let modules = m.children()
         if let child = modules["child"], let child = child.unwrap() as? QuantizedLinear {
             // make sure it was updated
-            XCTAssertEqual(child.shape.0, 256)
+            #expect(child.shape.0 == 256)
         } else {
-            XCTFail("child is nil or not QuantizedLinear")
+            Issue.record("child is nil or not QuantizedLinear")
         }
     }
 
-    func testUpdateModuleArray() throws {
+    @Test func updateModuleArray() throws {
         // update an array of modules
         let m = newStructureModule()
         var u = NestedDictionary<String, Module>()
@@ -527,31 +524,28 @@ class ModuleTests: XCTestCase {
         if let interior2 = modules["interior2"], let interior2 = interior2.unwrap() as? Module {
             let modules = interior2.children()
             if let children = modules["children"], let children = children.unwrap() as? [Linear] {
-                XCTAssertEqual(children.count, 3)
-                XCTAssertEqual(children[0].shape.0, 5)
+                #expect(children.count == 3)
+                #expect(children[0].shape.0 == 5)
             } else {
-                XCTFail("children is nil or wrong type")
+                Issue.record("children is nil or wrong type")
             }
         } else {
-            XCTFail("interior2 is nil")
+            Issue.record("interior2 is nil")
         }
     }
 
-    func testUpdateModuleNotWrapped() throws {
+    @Test func updateModuleNotWrapped() throws {
         // a module that doesn't have @ModuleInfo
         let m = newStructureModule()
         var u = NestedDictionary<String, Module>()
         u["nonWrappedChild"] = .value(Linear(4, 4))
 
-        do {
+        #expect(throws: UpdateError.self) {
             try m.update(modules: u, verify: .all)
-            XCTFail("should have thrown")
-        } catch {
-            print("Expected: \(error)")
         }
     }
 
-    func testLinearUpdateParametersAllSet() throws {
+    @Test func linearUpdateParametersAllSet() throws {
         class M: Module {
             let a: Linear
             let b: Linear
@@ -575,10 +569,10 @@ class ModuleTests: XCTestCase {
                 ])),
             verify: .all)
         let wA = m.a.weight.asArray(Int.self)
-        XCTAssertEqual(wA, [0, 1])
+        #expect(wA == [0, 1])
         let wB = m.b.weight.asArray(Int.self)
-        XCTAssertEqual(wB, Array(0 ..< 12))
-        XCTAssertThrowsError(
+        #expect(wB == Array(0 ..< 12))
+        let error = #expect(throws: UpdateError.self) {
             try m.update(
                 parameters: .init(
                     item: .dictionary([
@@ -586,23 +580,22 @@ class ModuleTests: XCTestCase {
                         "b": .dictionary(["weight": .value(weightsB)]),
                     ])),
                 verify: .all)
-        ) { error in
-            guard let error = error as? UpdateError,
-                case .keyNotFound(let path, let modules) = error
-            else {
-                XCTFail("Expected to fail with UpdateError.keyNotFound, but got: \(error)")
-                return
-            }
-            // should be a.bias or b.bias (random order as it is a dict)
-            XCTAssertEqual(path.last, "bias")
-            XCTAssertEqual(modules, ["M", "Linear"])
         }
+        guard case .keyNotFound(let path, let modules)? = error else {
+            Issue.record(
+                "Expected to fail with UpdateError.keyNotFound, but got: \(String(describing: error))"
+            )
+            return
+        }
+        // should be a.bias or b.bias (random order as it is a dict)
+        #expect(path.last == "bias")
+        #expect(modules == ["M", "Linear"])
     }
 
-    func testLinearUpdateParametersVerifyAll() throws {
+    @Test func linearUpdateParametersVerifyAll() throws {
         let linear = Linear(1, 2, bias: false)
 
-        XCTAssertEqual(linear.weight.shape, [2, 1])
+        #expect(linear.weight.shape == [2, 1])
 
         let correctWeights = MLXArray(0 ..< 2).reshaped([2, 1])
 
@@ -610,49 +603,48 @@ class ModuleTests: XCTestCase {
             parameters: .init(item: .dictionary(["weight": .value(correctWeights)])),
             verify: .all)
 
-        XCTAssertEqual(linear.weight.shape, [2, 1])
+        #expect(linear.weight.shape == [2, 1])
 
         let transposedWeights = MLXArray(0 ..< 2).reshaped([1, 2])
-        XCTAssertThrowsError(
+        let error = #expect(throws: UpdateError.self) {
             try linear.update(
                 parameters: .init(item: .dictionary(["weight": .value(transposedWeights)])),
                 verify: .all)
-        ) { error in
-            guard let error = error as? UpdateError,
-                case .mismatchedSize(
-                    let
-                        path, let modules, expectedShape: let expectedShape,
-                    actualShape: let actualShape) =
-                    error
-            else {
-                XCTFail("Expected to fail with UpdateError.mismatchedSize, but got: \(error)")
-                return
-            }
-            XCTAssertEqual(expectedShape, [2, 1])
-            XCTAssertEqual(actualShape, [1, 2])
-            XCTAssertEqual(path, ["weight"])
-            XCTAssertEqual(
-                error.errorDescription,
-                "Mismatched parameter weight in Linear shape. Actual [1, 2], expected [2, 1]")
         }
-    }
-
-    func testLinearUpdateParametersVerifyNone() throws {
+        guard
+            let error,
+            case .mismatchedSize(
+                let path, _, expectedShape: let expectedShape, actualShape: let actualShape) = error
+        else {
+            Issue.record(
+                "Expected to fail with UpdateError.mismatchedSize, but got: \(String(describing: error))"
+            )
+            return
+        }
+        #expect(expectedShape == [2, 1])
+        #expect(actualShape == [1, 2])
+        #expect(path == ["weight"])
+        #expect(
+            error.errorDescription
+                == "Mismatched parameter weight in Linear shape. Actual [1, 2], expected [2, 1]")
+    }
+
+    @Test func linearUpdateParametersVerifyNone() throws {
         let linear = Linear(1, 2, bias: false)
 
-        XCTAssertEqual(linear.weight.shape, [2, 1])
+        #expect(linear.weight.shape == [2, 1])
 
         let transposedWeights = MLXArray(0 ..< 2).reshaped([1, 2])
         try linear.update(
             parameters: .init(item: .dictionary(["weight": .value(transposedWeights)])),
             verify: .none)
 
-        XCTAssertEqual(
-            linear.weight.shape, [1, 2],
+        #expect(
+            linear.weight.shape == [1, 2],
             "In verify none mode, parameters can be updated with a different shape")
     }
 
-    func testDictionaryParameterUpdate() {
+    @Test func dictionaryParameterUpdate() {
         class Model: Module {
             var layers: [String: Linear] = ["a": Linear(2, 2)]
         }
@@ -664,10 +656,10 @@ class ModuleTests: XCTestCase {
 
         let weight = m.layers["a"]!.weight
         eval(weight)
-        XCTAssertEqual(weight[0, 0].item(Float.self), 1.0)
+        #expect(weight[0, 0].item(Float.self) == 1.0)
     }
 
-    func testQuantize() throws {
+    @Test func quantizeModel() throws {
         class C: Module {
             @ModuleInfo
             var child = Linear(256, 256)
@@ -681,11 +673,11 @@ class ModuleTests: XCTestCase {
         let m = M()
         quantize(model: m)
 
-        XCTAssertTrue(m.module.child is QuantizedLinear)
-        XCTAssertTrue(m.module.child.shape == (256, 256))
+        #expect(m.module.child is QuantizedLinear)
+        #expect(m.module.child.shape == (256, 256))
     }
 
-    func testAlreadyQuantized() throws {
+    @Test func alreadyQuantized() throws {
         // should not quantize an already quantize layer -- that would be silly
         class C: Module {
             @ModuleInfo
@@ -700,11 +692,11 @@ class ModuleTests: XCTestCase {
         let m = M()
         quantize(model: m)
 
-        XCTAssertTrue(m.module.child is QuantizedLinear)
-        XCTAssertTrue(m.module.child.shape == (256, 256))
+        #expect(m.module.child is QuantizedLinear)
+        #expect(m.module.child.shape == (256, 256))
     }
 
-    func testQuantizePredicate() throws {
+    @Test func quantizePredicate() throws {
         class C: Module {
             @ModuleInfo
             var child1 = Linear(256, 256)
@@ -729,11 +721,11 @@ class ModuleTests: XCTestCase {
                 }
             })
 
-        XCTAssertTrue(m.module.child1 is QuantizedLinear)
-        XCTAssertFalse(m.module.child2 is QuantizedLinear)
+        #expect(m.module.child1 is QuantizedLinear)
+        #expect(!(m.module.child2 is QuantizedLinear))
     }
 
-    func testNilParameters() {
+    @Test func nilParameters() {
         class M: Module {
             @ParameterInfo var a: MLXArray?
             let b: MLXArray?
@@ -747,10 +739,10 @@ class ModuleTests: XCTestCase {
         let m = M()
         let parameters = m.parameters()
 
-        XCTAssertTrue(parameters.isEmpty)
+        #expect(parameters.isEmpty)
     }
 
-    func testNilModules() {
+    @Test func nilModules() {
         class M: Module {
             @ModuleInfo var a: Linear?
             let b: Linear?
@@ -764,10 +756,10 @@ class ModuleTests: XCTestCase {
         let m = M()
         let children = m.children()
 
-        XCTAssertTrue(children.isEmpty)
+        #expect(children.isEmpty)
     }
 
-    func testBatchNormItems() {
+    @Test func batchNormItems() {
         // BatchNorm is interesting because:
         // - it has optional parameters
         // - it has optional parameters with keys
@@ -778,13 +770,12 @@ class ModuleTests: XCTestCase {
         n1.train(true)
 
         // it should have all the parameters present
-        XCTAssertEqual(
-            Set(n1.parameters().flattened().map { $0.0 }),
-            Set(["weight", "bias", "running_mean", "running_var"]))
+        #expect(
+            Set(n1.parameters().flattened().map { $0.0 })
+                == Set(["weight", "bias", "running_mean", "running_var"]))
 
         // only two of them are trainable
-        XCTAssertEqual(
-            Set(n1.trainableParameters().flattened().map { $0.0 }), Set(["weight", "bias"]))
+        #expect(Set(n1.trainableParameters().flattened().map { $0.0 }) == Set(["weight", "bias"]))
 
         _ = n1(MLXArray(0 ..< 4, [1, 1, 4]))
 
@@ -797,10 +788,10 @@ class ModuleTests: XCTestCase {
             let e: [Float] = [0, 0.1, 0.2, 0.3]
 
             for item in zip(a, e) {
-                XCTAssertEqual(item.0, item.1, accuracy: 0.01)
+                expectEqual(item.0, item.1, accuracy: 0.01)
             }
         default:
-            XCTFail("should be a .value(array)")
+            Issue.record("should be a .value(array)")
         }
 
         let v = parameters["running_var"]
@@ -810,14 +801,14 @@ class ModuleTests: XCTestCase {
             let e: [Float] = [0.9, 0.9, 0.9, 0.9]
 
             for item in zip(a, e) {
-                XCTAssertEqual(item.0, item.1, accuracy: 0.01)
+                expectEqual(item.0, item.1, accuracy: 0.01)
             }
         default:
-            XCTFail("should be a .value(array)")
+            Issue.record("should be a .value(array)")
         }
     }
 
-    func testUpdateModulesNil() {
+    @Test func updateModulesNil() {
         // test where a tuple is updated with a nil value, e.g. when
         // a triple with two Linear modules is quantized, the third value
         // will be nil but is not nullable.  The code should copy forward
@@ -839,11 +830,11 @@ class ModuleTests: XCTestCase {
 
         quantize(model: pm, groupSize: 64, bits: 8)
 
-        XCTAssertTrue(pm.mlp.0 is QuantizedLinear)
-        XCTAssertTrue(pm.mlp.2 is QuantizedLinear)
+        #expect(pm.mlp.0 is QuantizedLinear)
+        #expect(pm.mlp.2 is QuantizedLinear)
     }
 
-    func testCompositeLayer() throws {
+    @Test func compositeLayer() throws {
         // a test of making a LoRA layer as a composite -- it adapts
         // another layer but mixes their properties.  this is not
         // necessarily the way to implement LoRA (you can also use subclasses
@@ -926,17 +917,17 @@ class ModuleTests: XCTestCase {
         let linearProperties = linear.parameters()
         let loraProperties = lora.parameters()
 
-        XCTAssertEqual(linearProperties.count, 2)
-        XCTAssertEqual(loraProperties.count, 4)
+        #expect(linearProperties.count == 2)
+        #expect(loraProperties.count == 4)
 
         try lora.update(parameters: loraProperties.mapValues { $0 + 1 }, verify: .all)
 
         let trainable = lora.trainableParameters()
-        XCTAssertEqual(trainable.count, 2)
-        XCTAssertTrue(trainable["lora_a"] != nil)
+        #expect(trainable.count == 2)
+        #expect(trainable["lora_a"] != nil)
     }
 
-    func testModulesWithMLXArrayProperties() throws {
+    @Test func modulesWithMLXArrayProperties() throws {
         // https://github.com/ml-explore/mlx-swift-examples/issues/218
 
         class SuScaledRotaryEmbedding: Module {
@@ -950,16 +941,16 @@ class ModuleTests: XCTestCase {
         let rope = SuScaledRotaryEmbedding()
 
         // no parameters
-        XCTAssertEqual(rope.parameters().count, 0)
+        #expect(rope.parameters().count == 0)
 
         // but it can see the _freqs property
-        XCTAssertEqual(rope.items().count, 1)
+        #expect(rope.items().count == 1)
 
         // this should not throw because _freqs is not considered
         try rope.update(parameters: .init(), verify: .all)
     }
 
-    func testOptionalModuleUpdates() throws {
+    @Test func optionalModuleUpdates() throws {
         class M: Module {
             @ModuleInfo(key: "a") var a: Linear = Linear(1, 1)
             @ModuleInfo(key: "b") var b: Linear?
@@ -972,33 +963,27 @@ class ModuleTests: XCTestCase {
             // Test setting a value for an optional module
             let u = ModuleChildren(values: ["b": .value(b)])
             try m.update(modules: u, verify: .noUnusedKeys)
-            XCTAssertEqual(m.children().count, 2)
+            #expect(m.children().count == 2)
         }
 
         do {
             // Test resetting the optional module to nil
             let u = ModuleChildren(values: ["b": .none])
             try m.update(modules: u, verify: .noUnusedKeys)
-            XCTAssertEqual(m.children().count, 1)
+            #expect(m.children().count == 1)
         }
 
         do {
             // Test that setting nil on a non-optional module throws an error
             let u = ModuleChildren(values: ["a": .none])
-            XCTAssertThrowsError(
+            #expect(throws: UpdateError.self) {
                 try m.update(modules: u, verify: .noUnusedKeys)
-            ) { error in
-                switch error {
-                case _ as UpdateError:
-                    XCTAssertEqual(m.children().count, 1)
-                default:
-                    XCTFail("Wrong error type")
-                }
             }
+            #expect(m.children().count == 1)
         }
     }
 
-    func testSequentialQuantization() {
+    @Test func sequentialQuantization() {
         let s = Sequential {
             Linear(64, 64)
             Linear(64, 64)
diff --git a/Tests/MLXTests/NestedTests.swift b/Tests/MLXTests/NestedTests.swift
index 05eb1f0c..afaea073 100644
--- a/Tests/MLXTests/NestedTests.swift
+++ b/Tests/MLXTests/NestedTests.swift
@@ -2,12 +2,10 @@
 
 import Foundation
 import MLX
-import XCTest
+import Testing
 
-class NestedTests: XCTestCase {
-
-    override class func setUp() {
-    }
+@Suite
+struct NestedTests {
 
     static let defaultValues = [10, 1, 2, 1, 2, 3, 10, 20, 30]
 
@@ -29,45 +27,45 @@ class NestedTests: XCTestCase {
         return result
     }
 
-    func testCollection() {
+    @Test func collection() {
         let n = newNested()
-        XCTAssertEqual(n.count, 4)
-        XCTAssertEqual(Set(n.keys), Set(["value", "array1", "array2", "dictionary"]))
+        #expect(n.count == 4)
+        #expect(Set(n.keys) == Set(["value", "array1", "array2", "dictionary"]))
     }
 
-    func testMapValues() {
+    @Test func mapValues() {
         let n = newNested()
         let incremented = n.mapValues { $0 + 1 }
         let expected = newNested(values: Self.defaultValues.map { $0 + 1 })
 
-        XCTAssertEqual(incremented, expected)
+        #expect(incremented == expected)
     }
 
-    func testCompactMapValues() {
+    @Test func compactMapValues() {
         let n = newNested()
         let trimmed = n.compactMapValues { $0 % 2 == 1 ? $0 : nil }
-        XCTAssertEqual(trimmed.count, 2)
-        XCTAssertEqual(trimmed["array1"], .array([.value(1)]))
+        #expect(trimmed.count == 2)
+        #expect(trimmed["array1"] == .array([.value(1)]))
     }
 
-    func testFlatten1() {
+    @Test func flatten1() {
         let v: NestedItem<String, Int> = .array([.array([.array([.value(10)])])])
 
         let f = v.flattened()
-        XCTAssertEqual(f.description, #"[("0.0.0", 10)]"#)
+        #expect(f.description == #"[("0.0.0", 10)]"#)
 
         let uf = NestedItem<String, Int>.unflattened(f)
-        XCTAssertEqual(uf, v)
+        #expect(uf == v)
     }
 
-    func testFlatten() {
+    @Test func flatten() {
         let n = newNested()
         let f = n.flattened()
         let n2 = NestedDictionary<String, Int>.unflattened(f)
-        XCTAssertEqual(n, n2)
+        #expect(n == n2)
     }
 
-    func testFlattenedValues() {
+    @Test func flattenedValues() {
         // replacingValues() uses the original structure
         // and makes an identical structure with new values
         // from a flat array of values
@@ -77,10 +75,10 @@ class NestedTests: XCTestCase {
         let n2 = n.replacingValues(with: f)
 
         let expected = n.mapValues { $0 + 1 }
-        XCTAssertEqual(n2, expected)
+        #expect(n2 == expected)
     }
 
-    func testMap2() {
+    @Test func map2() {
         // map 2 parallel structures
         var d1 = NestedDictionary<String, Int>()
         d1["layers"] = .array([
@@ -100,10 +98,10 @@ class NestedTests: XCTestCase {
 
         let expected = d1.mapValues { "value = \($0)" }
 
-        XCTAssertEqual(r2, expected)
+        #expect(r2 == expected)
     }
 
-    func testMap3() {
+    @Test func map3() {
         // map 3 parallel structures
         var d1 = NestedDictionary<String, Int>()
         d1["layers"] = .array([
@@ -132,8 +130,8 @@ class NestedTests: XCTestCase {
         let expected2 = d2.mapValues { -$0 }
         let (_, expected3) = d1.mapValues(d2) { (0, "value = \($0) + \($1 ?? 0)") }
 
-        XCTAssertEqual(r2, expected2)
-        XCTAssertEqual(r3, expected3)
+        #expect(r2 == expected2)
+        #expect(r3 == expected3)
     }
 
 }
diff --git a/Tests/MLXTests/OptimizerTests.swift b/Tests/MLXTests/OptimizerTests.swift
index 08d7d87c..4393257b 100644
--- a/Tests/MLXTests/OptimizerTests.swift
+++ b/Tests/MLXTests/OptimizerTests.swift
@@ -62,100 +62,127 @@ class OptimizerTests: XCTestCase {
         }
     }
 
-    func checkTrain<T>(optimizer: OptimizerBase<T>, compile: Bool = false) {
-
-        // measure the distance from the prediction (model(x)) and the
-        // ground truth (y).  this gives feedback on how close the
-        // prediction is from matching the truth
-        @Sendable
-        func loss(model: LinearFunctionModel, x: MLXArray, y: MLXArray) -> MLXArray {
-            mseLoss(predictions: model(x), targets: y, reduction: .mean)
-        }
+    /// Train ``LinearFunctionModel`` with `optimizer` and check that the loss drops.
+    ///
+    /// This checks that the optimizer is wired up and moves the parameters the right
+    /// way, not how fast it converges -- `MLXIntegrationTests` compares the optimizer
+    /// math against python mlx.  Each optimizer gets settings that suit it (see the
+    /// callers): one learning rate does not fit them all, e.g. SGD's step scales with
+    /// the gradient while Lion's is always `±learningRate`.
+    ///
+    /// The settings and step count were chosen by simulating this training loop (the
+    /// update rules from `Optimizers.swift` over 2000 seeds): with 300 steps the worst
+    /// case for every optimizer keeps at most 2% of the starting loss, against the 10%
+    /// required here.
+    ///
+    /// - Parameters:
+    ///   - optimizer: the optimizer to train with
+    ///   - steps: number of training steps
+    ///   - compile: compile the training step (an example of compiling a step that
+    ///     updates the model and the optimizer)
+    func checkTrain<T>(
+        optimizer: OptimizerBase<T>, steps: Int = 300, compile: Bool = false,
+        file: StaticString = #filePath, line: UInt = #line
+    ) {
+        // the model initialization and the training data are random -- use a fixed
+        // seed so the result is repeatable.  the state is task local so tests running
+        // in parallel cannot draw from it (see RandomSeedTrait)
+        withRandomState(MLXRandom.RandomState(seed: 0)) {
+
+            // measure the distance from the prediction (model(x)) and the
+            // ground truth (y).  this gives feedback on how close the
+            // prediction is from matching the truth
+            @Sendable
+            func loss(model: LinearFunctionModel, x: MLXArray, y: MLXArray) -> MLXArray {
+                mseLoss(predictions: model(x), targets: y, reduction: .mean)
+            }
 
-        let model = LinearFunctionModel()
-        eval(model)
+            let model = LinearFunctionModel()
+            eval(model)
 
-        // the optimizer will use the gradients update the model parameters
-        let optimizer = SGD(learningRate: 1e-1)
+            // these are the target parameters
+            let m = 0.25
+            let b = 7
 
-        // these are the target parameters
-        let m = 0.25
-        let b = 7
+            // a fixed set of points to measure the loss before and after training --
+            // the training batches are random, so their losses are too noisy to compare
+            let xEval = MLX.linspace(-5.0, 5.0, count: 21).reshaped([21, 1])
+            let yEval = m * xEval + b
+            let initialLoss = loss(model: model, x: xEval, y: yEval).item(Float.self)
 
-        let lg = valueAndGrad(model: model, loss)
+            let lg = valueAndGrad(model: model, loss)
 
-        func step(_ x: MLXArray, _ y: MLXArray) -> MLXArray {
-            let (loss, grads) = lg(model, x, y)
-            optimizer.update(model: model, gradients: grads)
-            return loss
-        }
+            // the optimizer will use the gradients to update the model parameters
+            func step(_ x: MLXArray, _ y: MLXArray) -> MLXArray {
+                let (loss, grads) = lg(model, x, y)
+                optimizer.update(model: model, gradients: grads)
+                return loss
+            }
 
-        let resolvedStep =
-            compile
-            ? MLX.compile(inputs: [model, optimizer], outputs: [model, optimizer], step) : step
+            let resolvedStep =
+                compile
+                ? MLX.compile(inputs: [model, optimizer], outputs: [model, optimizer], step) : step
 
-        // run a number of epochs
-        var lastLoss: MLXArray!
-        for _ in 0 ..< 30 {
-            // print("target: b = \(b), m = \(m)")
-            // print("parameters: \(model.parameters())")
+            for _ in 0 ..< steps {
+                // generate random training data along with the ground truth.
+                // notice that the shape is [B, 1] where B is the batch
+                // dimension -- this allows us to train on 10 samples simultaneously
+                let x = MLXRandom.uniform(low: -5.0, high: 5.0, [10, 1])
+                let y = m * x + b
+                eval(x, y)
 
-            // generate random training data along with the ground truth.
-            // notice that the shape is [B, 1] where B is the batch
-            // dimension -- this allows us to train on 10 samples simultaneously
-            let x = MLXRandom.uniform(low: -5.0, high: 5.0, [10, 1])
-            let y = m * x + b
-            eval(x, y)
+                // compute the loss and gradients.  use the optimizer
+                // to adjust the parameters closer to the target
+                _ = resolvedStep(x, y)
 
-            // compute the loss and gradients.  use the optimizer
-            // to adjust the parameters closer to the target
-            let loss = resolvedStep(x, y)
+                eval(model, optimizer)
+            }
 
-            eval(model, optimizer)
+            let finalLoss = loss(model: model, x: xEval, y: yEval).item(Float.self)
 
-            lastLoss = loss
+            // the loss should drop by at least 90%
+            XCTAssertLessThan(
+                finalLoss, initialLoss * 0.1,
+                "\(type(of: optimizer)) (compile: \(compile)): loss \(initialLoss) -> \(finalLoss)",
+                file: file, line: line)
         }
-
-        // it should reach this loss
-        XCTAssertLessThan(lastLoss.item(Float.self), 0.1)
-
-        print("final loss: \(lastLoss!)")
     }
 
     // MARK: - Integration Tests
     //
     // integration tests:
     // - verify shapes match input (sort of a smoke test, copied from python)
-    // - verify that the otpimizer actually converges for a simple model
+    // - verify that training with the optimizer reduces the loss of a simple model
 
     func testSGD() {
         checkShape(optimizer: SGD(learningRate: 0.1))
         checkTrain(optimizer: SGD(learningRate: 0.1))
+
+        // the same training loop with the step compiled
         checkTrain(optimizer: SGD(learningRate: 0.1), compile: true)
     }
 
     func testRMSprop() {
         checkShape(optimizer: RMSprop(learningRate: 0.1))
         checkTrain(optimizer: RMSprop(learningRate: 0.1))
-        checkTrain(optimizer: RMSprop(learningRate: 0.1), compile: true)
     }
 
     func testAdaGrad() {
         checkShape(optimizer: AdaGrad(learningRate: 0.1))
-        checkTrain(optimizer: AdaGrad(learningRate: 0.1))
-        checkTrain(optimizer: AdaGrad(learningRate: 0.1), compile: true)
+        // the step shrinks as the squared gradients accumulate
+        checkTrain(optimizer: AdaGrad(learningRate: 1.0))
     }
 
     func testAdaDelta() {
         checkShape(optimizer: AdaDelta(learningRate: 0.1))
-        checkTrain(optimizer: AdaDelta(learningRate: 0.1))
-        checkTrain(optimizer: AdaDelta(learningRate: 0.1), compile: true)
+        // the step starts near sqrt(eps) and grows slowly -- with the default eps
+        // (1e-6) it barely moves in a few hundred steps, whatever the learning rate
+        checkTrain(optimizer: AdaDelta(learningRate: 1.0, eps: 1e-3))
     }
 
     func testAdam() {
         checkShape(optimizer: Adam(learningRate: 0.1))
         checkTrain(optimizer: Adam(learningRate: 0.1))
-        checkTrain(optimizer: Adam(learningRate: 0.1), compile: true)
     }
 
     func testAdamBiasCorrection() {
@@ -179,7 +206,6 @@ class OptimizerTests: XCTestCase {
     func testAdamW() {
         checkShape(optimizer: AdamW(learningRate: 0.1))
         checkTrain(optimizer: AdamW(learningRate: 0.1))
-        checkTrain(optimizer: AdamW(learningRate: 0.1), compile: true)
     }
 
     func testAdamWBiasCorrection() {
@@ -204,20 +230,21 @@ class OptimizerTests: XCTestCase {
 
     func testAdamax() {
         checkShape(optimizer: Adamax(learningRate: 0.1))
-        checkTrain(optimizer: Adamax(learningRate: 0.1))
-        checkTrain(optimizer: Adamax(learningRate: 0.1), compile: true)
+        checkTrain(optimizer: Adamax(learningRate: 1.0))
     }
 
     func testLion() {
         checkShape(optimizer: Lion(learningRate: 0.1))
+        // every step is exactly ±learningRate per parameter
         checkTrain(optimizer: Lion(learningRate: 0.1))
-        checkTrain(optimizer: Lion(learningRate: 0.1), compile: true)
     }
 
     func testAdafactor() {
         checkShape(optimizer: Adafactor(learningRate: 0.1))
-        checkTrain(optimizer: Adafactor(learningRate: 0.1))
-        checkTrain(optimizer: Adafactor(learningRate: 0.1), compile: true)
+        // by default (relativeStep) the learning rate is ignored and the step is
+        // relative to the size of the parameter, so parameters near zero hardly move
+        checkTrain(
+            optimizer: Adafactor(learningRate: 0.1, scaleParameter: false, relativeStep: false))
     }
 
     // MARK: - Muon
diff --git a/Tests/MLXTests/QuantizationTests.swift b/Tests/MLXTests/QuantizationTests.swift
index 357ac81e..73dc2cd2 100644
--- a/Tests/MLXTests/QuantizationTests.swift
+++ b/Tests/MLXTests/QuantizationTests.swift
@@ -6,26 +6,43 @@ import MLXNN
 import XCTest
 
 class QuantizationTests: XCTestCase {
+
+    func testQuantizationError() {
+        let a = MLXRandom.uniform(low: 0, high: 1, [23, 18])
+        do {
+            try withError {
+                let _ = quantized(a, groupSize: 123, bits: 7)
+                // note: at this point the return values from quantized
+                // are mlxNone and unusable
+            }
+            XCTFail("should have thrown an error")
+        } catch {
+            // expected
+        }
+    }
+
     func testQuantizedLinearAppliesNVFP4GlobalScaleOnMetal() {
         let (input, quantizedWeight, scales, biases, globalScale, bias, expected) =
             Device.withDefaultDevice(.cpu) {
-                MLXRandom.seed(7)
-                let weight = MLXRandom.uniform(low: -2, high: 2, [32, 32])
-                let input = MLXRandom.uniform(low: -1, high: 1, [4, 32])
-                let globalScale = MLX.max(abs(weight)).asType(.float32)
-                let bias = MLXRandom.uniform(low: -1, high: 1, [32])
-                let (quantizedWeight, scales, biases) = MLX.quantized(
-                    weight, groupSize: 16, bits: 4, mode: .nvfp4, globalScale: globalScale)
-                let expected =
-                    matmul(
-                        input,
-                        dequantized(
-                            quantizedWeight, scales: scales, biases: biases, groupSize: 16, bits: 4,
-                            mode: .nvfp4, globalScale: globalScale, dtype: .float32
-                        ).T)
-                    + bias
-                eval(quantizedWeight, scales, expected)
-                return (input, quantizedWeight, scales, biases, globalScale, bias, expected)
+                withRandomState(MLXRandom.RandomState(seed: 7)) {
+                    let weight = MLXRandom.uniform(low: -2, high: 2, [32, 32])
+                    let input = MLXRandom.uniform(low: -1, high: 1, [4, 32])
+                    let globalScale = MLX.max(abs(weight)).asType(.float32)
+                    let bias = MLXRandom.uniform(low: -1, high: 1, [32])
+                    let (quantizedWeight, scales, biases) = MLX.quantized(
+                        weight, groupSize: 16, bits: 4, mode: .nvfp4, globalScale: globalScale)
+                    let expected =
+                        matmul(
+                            input,
+                            dequantized(
+                                quantizedWeight, scales: scales, biases: biases, groupSize: 16,
+                                bits: 4,
+                                mode: .nvfp4, globalScale: globalScale, dtype: .float32
+                            ).T)
+                        + bias
+                    eval(quantizedWeight, scales, expected)
+                    return (input, quantizedWeight, scales, biases, globalScale, bias, expected)
+                }
             }
 
         let actual = Device.withDefaultDevice(.gpu) {
@@ -42,25 +59,26 @@ class QuantizationTests: XCTestCase {
             input, indices, quantizedWeight, scales, biases, globalScale, expectedLinear,
             expectedLookup
         ) = Device.withDefaultDevice(.cpu) {
-            MLXRandom.seed(11)
-            let weight = MLXRandom.uniform(low: -2, high: 2, [32, 32])
-            let input = MLXRandom.uniform(low: -1, high: 1, [4, 32])
-            let indices = MLXArray([0, 3, 7, 31])
-            let globalScale = MLX.max(abs(weight)).asType(.float32)
-            let (quantizedWeight, scales, biases) = MLX.quantized(
-                weight, groupSize: 16, bits: 4, mode: .nvfp4, globalScale: globalScale)
-            let dequantizedWeight = dequantized(
-                quantizedWeight, scales: scales, biases: biases, groupSize: 16, bits: 4,
-                mode: .nvfp4, globalScale: globalScale, dtype: .float32)
-            let expectedLinear = matmul(input, dequantizedWeight.T)
-            let expectedLookup = dequantized(
-                quantizedWeight, scales: scales, biases: biases, groupSize: 16, bits: 4,
-                mode: .nvfp4, globalScale: globalScale)[indices].asType(.bfloat16)
-            eval(quantizedWeight, scales, expectedLinear, expectedLookup)
-            return (
-                input, indices, quantizedWeight, scales, biases, globalScale, expectedLinear,
-                expectedLookup
-            )
+            withRandomState(MLXRandom.RandomState(seed: 11)) {
+                let weight = MLXRandom.uniform(low: -2, high: 2, [32, 32])
+                let input = MLXRandom.uniform(low: -1, high: 1, [4, 32])
+                let indices = MLXArray([0, 3, 7, 31])
+                let globalScale = MLX.max(abs(weight)).asType(.float32)
+                let (quantizedWeight, scales, biases) = MLX.quantized(
+                    weight, groupSize: 16, bits: 4, mode: .nvfp4, globalScale: globalScale)
+                let dequantizedWeight = dequantized(
+                    quantizedWeight, scales: scales, biases: biases, groupSize: 16, bits: 4,
+                    mode: .nvfp4, globalScale: globalScale, dtype: .float32)
+                let expectedLinear = matmul(input, dequantizedWeight.T)
+                let expectedLookup = dequantized(
+                    quantizedWeight, scales: scales, biases: biases, groupSize: 16, bits: 4,
+                    mode: .nvfp4, globalScale: globalScale)[indices].asType(.bfloat16)
+                eval(quantizedWeight, scales, expectedLinear, expectedLookup)
+                return (
+                    input, indices, quantizedWeight, scales, biases, globalScale, expectedLinear,
+                    expectedLookup
+                )
+            }
         }
 
         let (actualLinear, actualLookup) = Device.withDefaultDevice(.gpu) {
diff --git a/Tests/MLXTests/TestingUtils.swift b/Tests/MLXTests/TestingUtils.swift
new file mode 100644
index 00000000..7795439c
--- /dev/null
+++ b/Tests/MLXTests/TestingUtils.swift
@@ -0,0 +1,146 @@
+// Copyright © 2026 Apple Inc.
+
+import MLX
+import Testing
+
+// Swift Testing counterparts to the XCTest helpers in `Utils.swift`.
+//
+// `XCTAssert*` is not a reliable way to fail a Swift Testing `@Test`, so
+// tests converted to Swift Testing should use these instead.
+
+/// Prefix `detail` with the caller supplied `comment`, if any.
+private func message(_ comment: Comment?, _ detail: String) -> Comment {
+    if let comment {
+        Comment(rawValue: "\(comment.rawValue)\n\(detail)")
+    } else {
+        Comment(rawValue: detail)
+    }
+}
+
+/// Expect two arrays to have the same shape and contents.
+///
+/// Shape is checked explicitly because `allClose` broadcasts -- e.g. `[2]`
+/// would otherwise compare equal to `[[2, 2], [2, 2]]`.
+///
+/// - Returns: `true` if the expectation passed (an issue is recorded otherwise)
+@discardableResult
+func expectEqual(
+    _ array1: MLXArray, _ array2: MLXArray, rtol: Double = 1e-5, atol: Double = 1e-8,
+    _ comment: Comment? = nil, sourceLocation: SourceLocation = #_sourceLocation
+) -> Bool {
+    guard array1.shape == array2.shape else {
+        Issue.record(
+            message(
+                comment, "shapes differ: \(array1.shape) != \(array2.shape)\n\(array1)\n\(array2)"),
+            sourceLocation: sourceLocation)
+        return false
+    }
+    if !array1.allClose(array2, rtol: rtol, atol: atol).item(Bool.self) {
+        Issue.record(
+            message(comment, "contents differ:\n\(array1)\n\(array2)"),
+            sourceLocation: sourceLocation)
+        return false
+    }
+    return true
+}
+
+/// Expect two lists of arrays to be pairwise equal (see ``expectEqual(_:_:rtol:atol:_:sourceLocation:)``).
+@discardableResult
+func expectEqual(
+    _ array1: [MLXArray], _ array2: [MLXArray], rtol: Double = 1e-5, atol: Double = 1e-8,
+    _ comment: Comment? = nil, sourceLocation: SourceLocation = #_sourceLocation
+) -> Bool {
+    guard array1.count == array2.count else {
+        Issue.record(
+            message(comment, "counts differ: \(array1.count) != \(array2.count)"),
+            sourceLocation: sourceLocation)
+        return false
+    }
+    var result = true
+    for (e1, e2) in zip(array1, array2) {
+        result =
+            expectEqual(e1, e2, rtol: rtol, atol: atol, comment, sourceLocation: sourceLocation)
+            && result
+    }
+    return result
+}
+
+/// Expect two arrays to have the same shape but different contents.
+@discardableResult
+func expectNotEqual(
+    _ array1: MLXArray, _ array2: MLXArray, rtol: Double = 1e-5, atol: Double = 1e-8,
+    _ comment: Comment? = nil, sourceLocation: SourceLocation = #_sourceLocation
+) -> Bool {
+    guard array1.shape == array2.shape else {
+        Issue.record(
+            message(comment, "shapes differ: \(array1.shape) != \(array2.shape)"),
+            sourceLocation: sourceLocation)
+        return false
+    }
+    if array1.allClose(array2, rtol: rtol, atol: atol).item(Bool.self) {
+        Issue.record(
+            message(comment, "contents same:\n\(array1)\n\(array2)"),
+            sourceLocation: sourceLocation)
+        return false
+    }
+    return true
+}
+
+/// Expect two floating point scalars to be within `accuracy` of each other.
+///
+/// Swift Testing counterpart to `XCTAssertEqual(_:_:accuracy:)`.
+///
+/// - Returns: `true` if the expectation passed (an issue is recorded otherwise)
+@discardableResult
+func expectEqual<T: FloatingPoint>(
+    _ value1: T, _ value2: T, accuracy: T,
+    _ comment: Comment? = nil, sourceLocation: SourceLocation = #_sourceLocation
+) -> Bool {
+    // written so that NaN fails rather than passes
+    guard abs(value1 - value2) <= accuracy else {
+        Issue.record(
+            message(comment, "\(value1) is not within \(accuracy) of \(value2)"),
+            sourceLocation: sourceLocation)
+        return false
+    }
+    return true
+}
+
+// MARK: - Random state
+
+/// Run a test with a task-local random state seeded with `seed`, e.g.:
+///
+/// ```swift
+/// @Test(.randomSeed(7)) func somethingRandom() {
+///     // draws from the task-local state, as if MLXRandom.seed(7) had been called
+///     let a = MLXRandom.normal([2, 2])
+/// }
+/// ```
+///
+/// Use this instead of `MLXRandom.seed(_:)`.  The global random state is shared by
+/// every test in the process and Swift Testing runs tests in parallel, so another
+/// test can draw from it between the `seed` and the values a test depends on.
+///
+/// `MLXRandom.RandomState(seed: n)` starts from the same key as `MLXRandom.seed(n)`,
+/// so expected values computed with a global seed stay valid.  The state is a
+/// task local (see `withRandomState`) so it covers the whole test body, including
+/// child tasks and anything that draws implicitly, e.g. `Linear` initialization.
+struct RandomSeedTrait: TestTrait, TestScoping {
+    let seed: UInt64
+
+    func provideScope(
+        for test: Test, testCase: Test.Case?,
+        performing function: @Sendable () async throws -> Void
+    ) async throws {
+        try await withRandomState(MLXRandom.RandomState(seed: seed)) {
+            try await function()
+        }
+    }
+}
+
+extension Trait where Self == RandomSeedTrait {
+    /// See ``RandomSeedTrait``.
+    static func randomSeed(_ seed: UInt64) -> Self {
+        RandomSeedTrait(seed: seed)
+    }
+}
diff --git a/Tests/MLXTests/TransformTests.swift b/Tests/MLXTests/TransformTests.swift
index dc1c5598..57521867 100644
--- a/Tests/MLXTests/TransformTests.swift
+++ b/Tests/MLXTests/TransformTests.swift
@@ -71,43 +71,44 @@ class TransformTests: XCTestCase {
         // this isn't a real model but can exercise the
         // machinery
 
-        MLXRandom.seed(0)
+        withRandomState(MLXRandom.RandomState(seed: 0)) {
 
-        class M: Module, UnaryLayer {
-            let linear = Linear(5, 5)
+            class M: Module, UnaryLayer {
+                let linear = Linear(5, 5)
 
-            func callAsFunction(_ x: MLXArray) -> MLXArray {
-                relu(linear(x) * 5 + 1)
+                func callAsFunction(_ x: MLXArray) -> MLXArray {
+                    relu(linear(x) * 5 + 1)
+                }
             }
-        }
 
-        let m = M()
-        MLX.eval(m)
+            let m = M()
+            MLX.eval(m)
 
-        let x = MLXArray(0 ..< 5, [1, 5])
-        let y = MLXArray(0 ..< 5)
+            let x = MLXArray(0 ..< 5, [1, 5])
+            let y = MLXArray(0 ..< 5)
 
-        func loss(model: M, x: MLXArray, y: MLXArray) -> MLXArray {
-            crossEntropy(logits: model(x), targets: y, reduction: .mean)
-        }
+            func loss(model: M, x: MLXArray, y: MLXArray) -> MLXArray {
+                crossEntropy(logits: model(x), targets: y, reduction: .mean)
+            }
 
-        let lg = valueAndGrad(model: m, loss)
+            let lg = valueAndGrad(model: m, loss)
 
-        let (loss, grads) = lg(m, x, y)
+            let (loss, grads) = lg(m, x, y)
 
-        // we have known loss because of the random seed
-        XCTAssertEqual(loss.item(), Float(5.98534), accuracy: 0.001)
+            // we have known loss because of the random seed
+            XCTAssertEqual(loss.item(), Float(5.98534), accuracy: 0.001)
 
-        if let linear = grads["linear"] {
-            switch linear {
-            case .dictionary(let d):
-                XCTAssertNotNil(d["bias"])
-                XCTAssertNotNil(d["weight"])
-            default:
-                XCTFail("linear is not a dictionary")
+            if let linear = grads["linear"] {
+                switch linear {
+                case .dictionary(let d):
+                    XCTAssertNotNil(d["bias"])
+                    XCTAssertNotNil(d["weight"])
+                default:
+                    XCTFail("linear is not a dictionary")
+                }
+            } else {
+                XCTFail("missing linear key")
             }
-        } else {
-            XCTFail("missing linear key")
         }
     }
 
diff --git a/Tests/MLXTests/Utils.swift b/Tests/MLXTests/Utils.swift
index 00c829c8..4611d5ed 100644
--- a/Tests/MLXTests/Utils.swift
+++ b/Tests/MLXTests/Utils.swift
@@ -8,10 +8,12 @@ func assertEqual(
     _ array1: MLXArray, _ array2: MLXArray, rtol: Double = 1e-5, atol: Double = 1e-8,
     file: StaticString = #filePath, line: UInt = #line
 ) {
-    XCTAssertEqual(array1.shape, array2.shape, "shapes differ: \(array1.shape) != \(array2.shape)")
+    XCTAssertEqual(
+        array1.shape, array2.shape, "shapes differ: \(array1.shape) != \(array2.shape)",
+        file: file, line: line)
     XCTAssertTrue(
         array1.allClose(array2, rtol: rtol, atol: atol).item(Bool.self),
-        "contents differ:\n\(array1)\n\(array2)")
+        "contents differ:\n\(array1)\n\(array2)", file: file, line: line)
 }
 
 func assertEqual(
@@ -28,8 +30,10 @@ func assertNotEqual(
     _ array1: MLXArray, _ array2: MLXArray, rtol: Double = 1e-5, atol: Double = 1e-8,
     file: StaticString = #filePath, line: UInt = #line
 ) {
-    XCTAssertEqual(array1.shape, array2.shape, "shapes differ: \(array1.shape) != \(array2.shape)")
+    XCTAssertEqual(
+        array1.shape, array2.shape, "shapes differ: \(array1.shape) != \(array2.shape)",
+        file: file, line: line)
     XCTAssertFalse(
         array1.allClose(array2, rtol: rtol, atol: atol).item(Bool.self),
-        "contents same:\n\(array1)\n\(array2)")
+        "contents same:\n\(array1)\n\(array2)", file: file, line: line)
 }