replace and improve integration tests

.github/workflows/pull_request.yml

@@ -115,6 +115,7 @@ jobs:
         run: |
           xcrun xctest ~/Library/Developer/Xcode/DerivedData/mlx-swift-*/Build/Products/Debug/CmlxTests.xctest
           xcrun xctest ~/Library/Developer/Xcode/DerivedData/mlx-swift-*/Build/Products/Debug/MLXTests.xctest
+          xcrun xctest ~/Library/Developer/Xcode/DerivedData/mlx-swift-*/Build/Products/Debug/MLXIntegrationTests.xctest
 
       - name: Run Distributed Tests SwiftPM (Xcode, macOS)
         shell: sh

MAINTENANCE.md

@@ -91,12 +91,7 @@ dependencies: [.product(name: "MLX", package: "mlx-swift"),
                .product(name: "MLXOptimizers", package: "mlx-swift")]
 ```
 
-10. Update `tools/generate_integration_tests.py` as needed
-
-```
-import MLXNN
-@testable import MLXOptimizers
-```
+10. Update `tools/integration_tests/cases.py` as needed, regenerate tests if needed
 
 11. Update tests as needed
 

Package.swift

@@ -405,6 +405,12 @@ let package = Package(
                 "MLX", "MLXNN", "MLXOptimizers",
             ]
         ),
+        .testTarget(
+            name: "MLXIntegrationTests",
+            dependencies: [
+                "MLX", "MLXNN", "MLXOptimizers",
+            ]
+        ),
 
         // ------
         // Example programs

Source/MLX/Documentation.docc/Articles/converting-python.md

@@ -107,8 +107,8 @@ Note: some of the symbols are not linkable.
 `cos` | ``MLXArray/cos(stream:)``
 `cummax` | ``MLXArray/cummax(axis:reverse:inclusive:stream:)``
 `cummin` | ``MLXArray/cummin(axis:reverse:inclusive:stream:)``
-`cumprod` | ``MLXArray/cumprod(axis:reverse:inclusive:stream:)``
-`cumsum` | ``MLXArray/cumsum(axis:reverse:inclusive:stream:)``
+`cumprod` | ``MLXArray/cumprod(axis:reverse:inclusive:dtype:stream:)``
+`cumsum` | ``MLXArray/cumsum(axis:reverse:inclusive:dtype:stream:)``
 `exp` | ``MLXArray/exp(stream:)``
 `flatten` | ``MLXArray/flattened(start:end:stream:)``
 `log` | ``MLXArray/log(stream:)``
@@ -190,7 +190,7 @@ This is a mapping of `mx` free functions to their ``MLX`` counterparts.
 `identity` | ``MLXArray/identity(_:type:stream:)``
 `less` | ``MLX/less(_:_:stream:)``
 `less_equal` | ``MLX/lessEqual(_:_:stream:)``
-`linspace` | ``MLXArray/linspace(_:_:count:stream:)-(Int,Int,Int,StreamOrDevice)``
+`linspace` | ``MLXArray/linspace(_:_:count:dtype:stream:)-(Int,Int,Int,DType?,StreamOrDevice)``
 `load` | ``MLX/loadArray(url:stream:)`` and ``MLX/loadArrays(url:stream:)``
 `log` | ``MLX/log(_:stream:)``
 `log10` | ``MLX/log10(_:stream:)``
@@ -210,7 +210,7 @@ This is a mapping of `mx` free functions to their ``MLX`` counterparts.
 `negative` | ``MLX/negative(_:stream:)``
 `not_equal` | ``MLX/notEqual(_:_:stream:)``
 `ones` | ``MLXArray/ones(_:type:stream:)``
-`ones_like` | ``MLXArray/ones(like:stream:)``
+`ones_like` | ``MLXArray/ones(like:dtype:stream:)``
 `pad` | ``MLX/padded(_:width:mode:value:stream:)``
 `partition` | ``MLX/partitioned(_:kth:axis:stream:)``
 `power` | ``MLX/pow(_:_:stream:)-(MLXArray,MLXArray,_)``
@@ -255,4 +255,4 @@ This is a mapping of `mx` free functions to their ``MLX`` counterparts.
 `var` | ``MLX/variance(_:axes:keepDims:ddof:stream:)``
 `where` | ``MLX/which(_:_:_:stream:)``
 `zeros` | ``MLXArray/zeros(_:type:stream:)``
-`zeros_like` | ``MLXArray/zeros(like:stream:)``
+`zeros_like` | ``MLXArray/zeros(like:dtype:stream:)``

Source/MLX/Documentation.docc/Organization/cumulative.md

@@ -25,9 +25,9 @@ These are available as both methods on `MLXArray` and free functions.  They each
 - ``MLXArray/cummax(reverse:inclusive:stream:)``
 - ``MLXArray/cummin(axis:reverse:inclusive:stream:)``
 - ``MLXArray/cummin(reverse:inclusive:stream:)``
-- ``MLXArray/cumprod(axis:reverse:inclusive:stream:)``
+- ``MLXArray/cumprod(axis:reverse:inclusive:dtype:stream:)``
 - ``MLXArray/cumprod(reverse:inclusive:stream:)``
-- ``MLXArray/cumsum(axis:reverse:inclusive:stream:)``
+- ``MLXArray/cumsum(axis:reverse:inclusive:dtype:stream:)``
 - ``MLXArray/cumsum(reverse:inclusive:stream:)``
 - ``MLXArray/logCumsumExp(axis:reverse:inclusive:stream:)``
 - ``MLXArray/logCumsumExp(reverse:inclusive:stream:)``

Source/MLX/Documentation.docc/Organization/initialization.md

@@ -264,8 +264,8 @@ there are specific initializers to request it:
 - ``MLXArray/full(_:values:type:stream:)``
 - ``MLXArray/full(_:values:stream:)``
 - ``MLXArray/identity(_:type:stream:)``
-- ``MLXArray/linspace(_:_:count:stream:)-(Int,Int,Int,StreamOrDevice)``
-- ``MLXArray/linspace(_:_:count:stream:)-(Double,Double,Int,StreamOrDevice)``
+- ``MLXArray/linspace(_:_:count:dtype:stream:)-(Int,Int,Int,DType?,StreamOrDevice)``
+- ``MLXArray/linspace(_:_:count:dtype:stream:)-(Double,Double,Int,DType?,StreamOrDevice)``
 - ``MLXArray/repeated(_:count:axis:stream:)``
 - ``MLXArray/repeated(_:count:stream:)``
 - ``MLXArray/repeat(_:count:axis:stream:)``
@@ -282,8 +282,8 @@ there are specific initializers to request it:
 - ``MLX/full(_:values:type:stream:)``
 - ``MLX/full(_:values:stream:)``
 - ``MLX/identity(_:type:stream:)``
-- ``linspace(_:_:count:endpoint:stream:)-47tc1``
-- ``linspace(_:_:count:endpoint:stream:)-38sfd``
+- ``linspace(_:_:count:endpoint:dtype:stream:)-2b6eu``
+- ``linspace(_:_:count:endpoint:dtype:stream:)-8k1d2``
 - ``MLXArray/repeated(_:count:axis:stream:)``
 - ``MLXArray/repeated(_:count:stream:)``
 - ``MLX/repeat(_:count:axis:stream:)``

Source/MLX/Documentation.docc/free-functions.md

@@ -106,8 +106,8 @@ operations as methods for convenience.
 - ``MLX/full(_:values:type:stream:)``
 - ``MLX/full(_:values:stream:)``
 - ``MLX/identity(_:type:stream:)``
-- ``linspace(_:_:count:endpoint:stream:)-47tc1``
-- ``linspace(_:_:count:endpoint:stream:)-38sfd``
+- ``linspace(_:_:count:endpoint:dtype:stream:)-2b6eu``
+- ``linspace(_:_:count:endpoint:dtype:stream:)-8k1d2``
 - ``MLX/repeated(_:count:axis:stream:)``
 - ``MLX/repeated(_:count:stream:)``
 - ``MLX/repeat(_:count:axis:stream:)``

Source/MLX/Factory.swift

@@ -331,31 +331,34 @@ extension MLXArray {
         MLX.identity(n, dtype: dtype, stream: stream)
     }
 
-    /// Generate `num` evenly spaced numbers over interval `[start, stop]` for `BinaryInteger`.
+    /// Generate `count` evenly spaced numbers over interval `[start, stop]` for `BinaryInteger`.
     ///
-    /// Example:
+    /// The result is floating point (`float32` by default) even for integer
+    /// bounds -- see ``linspace(_:_:count:endpoint:dtype:stream:)-2b6eu``.
     ///
     /// ```swift
-    /// // Create a 50 element 1-D array with values from 0 to 50
-    /// let r = MLXArray.linspace(0, 50)
+    /// // [0, 0.5, 1] as float32
+    /// let r = MLXArray.linspace(0, 1, count: 3)
     /// ```
     ///
     /// - Parameters:
     ///     - start: start value
     ///     - stop: stop value
     ///     - count: number of samples
+    ///     - dtype: dtype of the result, `float32` if not specified
     ///     - stream: stream or device to evaluate on
     ///
     /// ### See Also
     /// - <doc:initialization>
-    /// - ``linspace(_:_:count:stream:)-92x6l``
+    /// - ``linspace(_:_:count:dtype:stream:)-3fx01``
     static public func linspace<T: HasDType>(
-        _ start: T, _ stop: T, count: Int = 50, stream: StreamOrDevice = .default
+        _ start: T, _ stop: T, count: Int = 50, dtype: DType? = nil,
+        stream: StreamOrDevice = .default
     ) -> MLXArray where T: BinaryInteger {
-        MLX.linspace(start, stop, count: count, stream: stream)
+        MLX.linspace(start, stop, count: count, dtype: dtype, stream: stream)
     }
 
-    /// Generate `num` evenly spaced numbers over interval `[start, stop]` for `BinaryFloatingPoint`.
+    /// Generate `count` evenly spaced numbers over interval `[start, stop]` for `BinaryFloatingPoint`.
     ///
     /// Example:
     ///
@@ -368,15 +371,17 @@ extension MLXArray {
     ///     - start: start value
     ///     - stop: stop value
     ///     - count: number of samples
+    ///     - dtype: dtype of the result, derived from `T` if not specified
     ///     - stream: stream or device to evaluate on
     ///
     /// ### See Also
     /// - <doc:initialization>
-    /// - ``linspace(_:_:count:stream:)-7m7eg``
+    /// - ``linspace(_:_:count:dtype:stream:)-9yqai``
     static public func linspace<T: HasDType>(
-        _ start: T, _ stop: T, count: Int = 50, stream: StreamOrDevice = .default
+        _ start: T, _ stop: T, count: Int = 50, dtype: DType? = nil,
+        stream: StreamOrDevice = .default
     ) -> MLXArray where T: BinaryFloatingPoint {
-        MLX.linspace(start, stop, count: count, stream: stream)
+        MLX.linspace(start, stop, count: count, dtype: dtype, stream: stream)
     }
 
     /// Generate values in the half-open interval `[0, stop)`.
@@ -1007,39 +1012,50 @@ public func identity(_ n: Int, dtype: DType, stream: StreamOrDevice = .default)
     return MLXArray(result)
 }
 
-/// Generate `num` evenly spaced numbers over interval `[start, stop]`.
+/// Generate `count` evenly spaced numbers over interval `[start, stop]`.
 ///
-/// Example:
+/// The result is floating point (`float32` by default), matching python's
+/// `mx.linspace`, even for integer bounds: evenly spaced values between two
+/// integers are generally not integers.  Pass `dtype:` for anything else:
 ///
 /// ```swift
-/// // Create a 50 element 1-D array with values from 0 to 50
-/// let r = MLXArray.linspace(0, 50)
+/// // [0, 0.5, 1] as float32 -- *not* [0, 0, 1] as int64
+/// let r = MLXArray.linspace(0, 1, count: 3)
+///
+/// // opt in to an integer (truncating) result
+/// let i = MLXArray.linspace(0, 10, count: 6, dtype: .int32)
 /// ```
 ///
 /// - Parameters:
 ///     - start: start value
 ///     - stop: stop value
 ///     - count: number of samples
 ///     - endpoint: if `true` then the endpoint is the last sample, if `false` it is a half-open interval
+///     - dtype: dtype of the result, `float32` if not specified
 ///     - stream: stream or device to evaluate on
 ///
 /// ### See Also
 /// - <doc:initialization>
-/// - ``linspace(_:_:count:endpoint:stream:)``
+/// - ``linspace(_:_:count:endpoint:dtype:stream:)-2b6eu``
 public func linspace<T: HasDType>(
     _ start: T, _ stop: T, count: Int = 50,
     endpoint: Bool = true,
+    dtype: DType? = nil,
     stream: StreamOrDevice = .default
 ) -> MLXArray where T: BinaryInteger {
     var result = mlx_array_new()
     mlx_linspace_endpoint(
-        &result, Double(start), Double(stop), count.int32, endpoint, T.dtype.cmlxDtype, stream.ctx)
+        &result, Double(start), Double(stop), count.int32, endpoint,
+        (dtype ?? .float32).cmlxDtype, stream.ctx)
     return MLXArray(result)
 }
 
-/// Generate `num` evenly spaced numbers over interval `[start, stop]`.
+/// Generate `count` evenly spaced numbers over interval `[start, stop]`.
 ///
-/// Example:
+/// The result dtype follows the bounds (`Float` -> `float32`,
+/// `Float16` -> `float16`) except that `Double` produces `float32`, matching
+/// ``MLXArray/init(_:)`` and python's `mx.linspace` -- `float64` is not
+/// supported on the GPU.  Pass `dtype:` to be explicit.
 ///
 /// ```swift
 /// // Create a 50 element 1-D array with values from 0 to 1
@@ -1051,19 +1067,24 @@ public func linspace<T: HasDType>(
 ///     - stop: stop value
 ///     - count: number of samples
 ///     - endpoint: if `true` then the endpoint is the last sample, if `false` it is a half-open interval
+///     - dtype: dtype of the result, derived from `T` if not specified
 ///     - stream: stream or device to evaluate on
 ///
 /// ### See Also
 /// - <doc:initialization>
-/// - ``linspace(_:_:count:endpoint:stream:)``
+/// - ``linspace(_:_:count:endpoint:dtype:stream:)-8k1d2``
 public func linspace<T: HasDType>(
     _ start: T, _ stop: T, count: Int = 50,
     endpoint: Bool = true,
+    dtype: DType? = nil,
     stream: StreamOrDevice = .default
 ) -> MLXArray where T: BinaryFloatingPoint {
+    // Double.dtype is float64, but we do not automatically promote to float64
+    // (see HasDType conformances) and it is not available on the GPU
+    let resolved = dtype ?? (T.dtype == .float64 ? .float32 : T.dtype)
     var result = mlx_array_new()
     mlx_linspace_endpoint(
-        &result, Double(start), Double(stop), count.int32, endpoint, T.dtype.cmlxDtype, stream.ctx)
+        &result, Double(start), Double(stop), count.int32, endpoint, resolved.cmlxDtype, stream.ctx)
     return MLXArray(result)
 }
 

Source/MLX/Ops+Array.swift

@@ -605,7 +605,7 @@ public func cummin(
 /// ### See Also
 /// - <doc:cumulative>
 /// - ``cumprod(_:reverse:inclusive:dtype:stream:)``
-/// - ``MLXArray/cumprod(axis:reverse:inclusive:stream:)``
+/// - ``MLXArray/cumprod(axis:reverse:inclusive:dtype:stream:)``
 public func cumprod(
     _ array: MLXArray, axis: Int, reverse: Bool = false, inclusive: Bool = true,
     dtype: DType? = nil,
@@ -629,7 +629,7 @@ public func cumprod(
 /// ### See Also
 /// - <doc:cumulative>
 /// - ``cumprod(_:axis:reverse:inclusive:dtype:stream:)``
-/// - ``MLXArray/cumprod(axis:reverse:inclusive:stream:)``
+/// - ``MLXArray/cumprod(axis:reverse:inclusive:dtype:stream:)``
 public func cumprod(
     _ array: MLXArray, reverse: Bool = false, inclusive: Bool = true,
     dtype: DType? = nil,
@@ -653,7 +653,7 @@ public func cumprod(
 /// ### See Also
 /// - <doc:cumulative>
 /// - ``cumsum(_:reverse:inclusive:dtype:stream:)``
-/// - ``MLXArray/cumsum(axis:reverse:inclusive:stream:)``
+/// - ``MLXArray/cumsum(axis:reverse:inclusive:dtype:stream:)``
 public func cumsum(
     _ array: MLXArray, axis: Int, reverse: Bool = false, inclusive: Bool = true,
     dtype: DType? = nil,
@@ -677,7 +677,7 @@ public func cumsum(
 /// ### See Also
 /// - <doc:cumulative>
 /// - ``cumsum(_:axis:reverse:inclusive:dtype:stream:)``
-/// - ``MLXArray/cumsum(axis:reverse:inclusive:stream:)``
+/// - ``MLXArray/cumsum(axis:reverse:inclusive:dtype:stream:)``
 public func cumsum(
     _ array: MLXArray, reverse: Bool = false, inclusive: Bool = true,
     dtype: DType? = nil,

Source/MLX/Ops.swift

@@ -1002,6 +1002,9 @@ public func convolve(
         if weightSize % 2 == 1 {
             padding = weightSize / 2
         } else {
+            // even sized weights use asymmetric padding -- this must match
+            // python's `mx.convolve()` so that the result is the centered
+            // `input.size` window of the full convolution
             let padLeft = weightSize / 2
             let padRight = max(0, padLeft - 1)
 
@@ -2184,7 +2187,7 @@ public func multiply(
 /// - <doc:arithmetic>
 public func nanToNum(
     _ array: MLXArray,
-    nan: Float = 0, posInf: Float? = 0, negInf: Float? = 0,
+    nan: Float = 0, posInf: Float? = nil, negInf: Float? = nil,
     stream: StreamOrDevice = .default
 ) -> MLXArray {
     let posInf = mlx_optional_float(value: posInf ?? 0, has_value: posInf != nil)
@@ -3110,15 +3113,15 @@ public func tanh(_ array: MLXArray, stream: StreamOrDevice = .default) -> MLXArr
 /// - Parameters:
 ///   - a: input array
 ///   - b: input array
-///   - axes: sum over the last `axes` dimensions
+///   - axes: sum over the last `axes` dimensions of `a` and the first `axes` of `b`
 ///   - stream: stream or device to evaluate on
 /// - Returns: tensor dot product
 ///
 /// ### See Also
 /// - <doc:arithmetic>
 /// - ``tensordot(_:_:axes:stream:)-(MLXArray,MLXArray,Int,StreamOrDevice)``
 public func tensordot(
-    _ a: MLXArray, _ b: MLXArray, axes: Int = 1, stream: StreamOrDevice = .default
+    _ a: MLXArray, _ b: MLXArray, axes: Int = 2, stream: StreamOrDevice = .default
 ) -> MLXArray {
     var result = mlx_array_new()
     mlx_tensordot_axis(&result, a.ctx, b.ctx, axes.int32, stream.ctx)

Source/MLXNN/Pooling.swift

@@ -41,9 +41,10 @@ open class Pool: Module, UnaryLayer {
 
         // Apply padding if any padding value is greater than 0
         if padding.contains(where: { $0 > 0 }) {
-            // batch and channel dimension get no padding
-            let padWidths: [IntOrPair] =
-                [0, 0] + padding.map { .init($0) } + [0, 0]
+            // one width per dimension: batch and channel get no padding, each
+            // spatial dimension is padded symmetrically (`IntOrPair` is already
+            // the (before, after) pair for a single dimension)
+            let padWidths: [IntOrPair] = [0] + padding.map { IntOrPair($0) } + [0]
             let paddingValue = paddingValue.asMLXArray(dtype: input.dtype)
             input = padded(input, widths: padWidths, mode: .constant, value: paddingValue)
         }

Source/MLXNN/PositionalEncoding.swift

@@ -156,19 +156,41 @@ final public class ALiBi: Module {
     public override init() {
     }
 
+    /// The per-head slopes, matching python's `ALiBi.create_alibi_slope()`.
+    ///
+    /// For a power of two head count the slopes are `2^(-8i/n)` for `i` in
+    /// `1...n`.  Otherwise python uses the slopes of the next power of two *below*
+    /// `n` and pads them with every other slope of the next power of two *above*,
+    /// which is not the same as extending the geometric series.
     static func alibiSlope(numHeads: Int) -> MLXArray {
-        let x = pow(pow(2, 8), (1 / Float(numHeads)))
-        let out = pow(x, -MLXArray(1 ..< (numHeads + 1)))
-        return out.expandedDimensions(axes: [-1, -2])
+        func slopes(_ n: Int) -> [Float] {
+            let log2n = Foundation.log2(Double(n))
+            if log2n == log2n.rounded(.down) {
+                let start = Foundation.pow(2.0, -Foundation.pow(2.0, 3 - log2n))
+                return (1 ... n).map { Float(Foundation.pow(start, Double($0))) }
+            }
+
+            let closestPowerOf2 = Int(Foundation.pow(2.0, log2n.rounded(.down)))
+            let interleaved = slopes(2 * closestPowerOf2)
+                .enumerated()
+                .filter { $0.offset.isMultiple(of: 2) }
+                .map { $0.element }
+                .prefix(n - closestPowerOf2)
+            return slopes(closestPowerOf2) + interleaved
+        }
+
+        return MLXArray(slopes(numHeads)).expandedDimensions(axes: [-1, -2])
     }
 
     static func alibiMatrix(key: Key) -> MLXArray {
         if let value = cache[key] {
             return value
         }
 
+        // x1 is a column and x2 a row so that the difference is the (q, k)
+        // distance matrix -- python: `x1[:, None] - x2[None, :]`
         let x1 = MLXArray(key.offset ..< key.qSequenceLength).expandedDimensions(axis: 1)
-        let x2 = MLXArray(0 ..< key.kSequenceLength).expandedDimensions(axis: 1)
+        let x2 = MLXArray(0 ..< key.kSequenceLength).expandedDimensions(axis: 0)
         let distanceMatrix = -abs(expandedDimensions((x1 - x2), axes: [0, 1]))
 
         let slope = alibiSlope(numHeads: key.numHeads)