replace and improve integration tests by davidkoski · Pull Request #477 · ml-explore/mlx-swift

davidkoski · 2026-09-11T18:07:17Z

Proposed changes

previously we used a python script to generate swift code and check output against known values from mlx (python)
this is replaced with a new and more powerful script (claude rewrite)
much more integration test coverage and tools to examine that
fixes per problems found by new integration tests
linspace output type, add dtype override
linspace default output type is float32 (matches python)
convolve .same padding side matches python (padRight)
nanToNum default for +/- infinity honors optionals (mlx defaults)
tensorDot axes default 1 -> 2 (matches python)
MLXNN Pooling padWidths matches python
MLXNN Alibi slope function matches python
MLXNN GRU apply bias to match python
MLXNN Lion optimizer default beta matches python
MLXNN Adafactor moving average matches python
MLXNN Muon zero power newton shulz matches python

Overall the number of tests increased like this (roughly):

area	total	integration	unit-only	none	no-swift	n/a
mx	240	66 -> 175	56 -> 16	80 -> 11	27	11
mx.fft	16	12 -> 16	2 -> 0	2 -> 0	0	0
mx.linalg	19	0 -> 11	3 -> 2	16 -> 6	0	0
mx.random	13	4 -> 11	6 -> 2	3 -> 0	0	0
mx.fast	7	1 -> 4	2 -> 1	2 -> 0	2	0
nn.layer	72	30 -> 56	9 -> 5	23 -> 1	10	0
nn.function	35	3 -> 29	4 -> 0	22 -> 0	6	0
nn.loss	14	0 -> 12	2 -> 0	10 -> 0	2	0
nn.init	10	2 -> 3	0	1 -> 0	7	0
optim	17	9 -> 16	7 -> 1	1 -> 0	0	0
swift-only	30	1	13 -> 12	16 -> 17	0	0

Checklist

Put an x in the boxes that apply.

I have read the CONTRIBUTING document
I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
I have added tests that prove my fix is effective or that my feature works
I have updated the necessary documentation (if needed)

AI disclosure: claude was used to write the new integration test generator (which in turn wrote the new integration tests), updated existing unit tests per issues found in integration tests, and in some cases make the code conform to the python reference.

- previously we used a python script to generate swift code and check output against known values from mlx (python) - this is replaced with a new and more powerful script (claude rewrite) - much more integration test coverage and tools to examine that - fixes per problems found by new integration tests - linspace output type, add dtype override - linspace default output type is float32 (matches python) - convolve .same padding side matches python (padRight) - nanToNum default for +/- infinity honors optionals (mlx defaults) - tensorDot axes default 1 -> 2 (matches python) - MLXNN Pooling padWidths matches python - MLXNN Alibi slop function matches python - MLXNN GRU apply bias to match python - MLXNN Lion optimizer default beta matches python - MLXNN Adafactor moving average matches python - MLXNN Muon zero power newton shulz matches python

davidkoski · 2026-09-11T18:08:55Z

-    /// - ``linspace(_:_:count:stream:)-92x6l``
+    /// - ``linspace(_:_:count:dtype:stream:)-92x6l``
    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 {


linspace over integers would produce integer result in swift, float32 in python

davidkoski · 2026-09-11T18:09:15Z

+            // `input.size` window of the full convolution
            let padLeft = weightSize / 2
-            let padRight = max(0, padLeft / 2 - 1)
+            let padRight = max(0, padLeft - 1)


convolve .same padding side matches python

davidkoski · 2026-09-11T18:09:45Z

 public func nanToNum(
    _ array: MLXArray,
-    nan: Float = 0, posInf: Float? = 0, negInf: Float? = 0,
+    nan: Float = 0, posInf: Float? = nil, negInf: Float? = nil,


use mlx default values for nanToNum +/- infinity

davidkoski · 2026-09-11T18:10:02Z

 /// - ``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


tensorDot axes default 1 -> 2 (matches python)

davidkoski · 2026-09-11T18:10:12Z

+            // 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]


MLXNN Pooling padWidths matches python

davidkoski · 2026-09-11T18:10:23Z

+    /// `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 {


MLXNN Alibi slope function matches python

davidkoski · 2026-09-11T18:10:39Z

+            } else if let bhn {
+                // matches python: without an incoming hidden state the hidden
+                // bias is still applied, gated by r
+                n = n + r * bhn


MLXNN GRU apply bias to match python

davidkoski · 2026-09-11T18:11:02Z

-    /// The coefficients used for computing running averages of the gradient and its square
-    public var betas: (Float, Float) = (0.9, 0.999)
+    /// The coefficients used for computing the gradient momentum and update direction
+    public var betas: (Float, Float) = (0.9, 0.99)


MLXNN Lion optimizer default beta matches python

davidkoski · 2026-09-11T18:11:12Z

-        return matmul(rFactor.expandedDimensions(axis: -1), cFactor.expandedDimensions(axis: 0))
+        // broadcast rather than matmul so this also works for parameters with more
+        // than two dimensions
+        return rFactor.expandedDimensions(axis: -1) * cFactor.expandedDimensions(axis: -2)


MLXNN Adafactor moving average matches python

davidkoski · 2026-09-11T18:11:23Z


    /// Orthogonalize a 2D matrix via a quintic Newton-Schulz iteration.
-    private func zeropowerViaNewtonSchulz5(_ input: MLXArray, steps: Int) -> MLXArray {
+    func zeropowerViaNewtonSchulz5(_ input: MLXArray, steps: Int) -> MLXArray {


MLXNN Muon zero power newton shulz matches python

davidkoski · 2026-09-11T18:11:59Z

@@ -0,0 +1,823 @@
+// Copyright © 2026 Apple Inc.
+//
+// GENERATED by tools/integration_tests -- DO NOT EDIT.


These are the new generated integration tests, replacing IntegrationTests.py (which was also generated)

davidkoski · 2026-09-11T18:12:11Z

Removed -- obsolete.

davidkoski · 2026-09-11T18:13:04Z

+inputs_case("normal/2d", {"a": Normal(S)})
+inputs_case("normal/2d/seed2", {"a": Normal(S)})
+inputs_case("normal/1d", {"a": Normal((17,))})


This replaces the generate_integration_tests.py

…them with TF32 disabled so we can guarantee consistent results on NAX vs non NAX hardware

nwh

LGTM -- my Claude reviewer noted an inconsistency in the README for the generator. Probably ok to remove the specific counts.

Co-authored-by: Nick Henderson <525247+nwh@users.noreply.github.com>

davidkoski requested review from ctymoszek and thechriswebb September 11, 2026 18:07

davidkoski commented Sep 11, 2026

View reviewed changes

Comment thread Tests/MLXTests/IntegrationTests.swift Outdated

davidkoski Sep 11, 2026

Copy link
Copy Markdown

Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Removed -- obsolete.

davidkoski commented Sep 11, 2026

View reviewed changes

davidkoski added 2 commits September 11, 2026 11:37

fix linspace disambiguation

eebe3ee

split integration tests into their own test target -- we need to run …

f1d3245

…them with TF32 disabled so we can guarantee consistent results on NAX vs non NAX hardware

This was referenced Sep 14, 2026

Add countNonzero to the Swift API #479

Merged

Release soon? ml-explore/mlx-swift-lm#617

Open

nwh self-requested a review September 16, 2026 20:56

nwh approved these changes Sep 16, 2026

View reviewed changes

Comment thread tools/integration_tests/README.md Outdated

Update tools/integration_tests/README.md

616b062

Co-authored-by: Nick Henderson <525247+nwh@users.noreply.github.com>

Uh oh!

Conversation

davidkoski commented Sep 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Proposed changes

Checklist

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

nwh left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

davidkoski commented Sep 11, 2026 •

edited

Loading