feat(gpu): MLX-batched pairs bootstrap with bit-identical Philox streams

- mlx-swift dependency in ZQGPU (the only module allowed to import MLX)
- Philox4x32-10 as vectorized MLX uint64 ops with the reference key
  schedule: GPU resample indices are bit-identical to the CPU generator
  (asserted per replicate) — the §5 reproducibility contract
- batched gather + float32 cross-products on GPU, float64 Cholesky solves
  on CPU (§5 hybrid); chunked to the memory budget
- conditioning: predictors standardized by fixed full-sample mean/scale
  (exact reparametrization) before forming X'X — raw data with an
  intercept lost 2-3 digits in float32
- MLX workaround: batched GEMM mis-accumulates small k x k outputs at
  batch >= 2 (~6e-4 rel, stride/fusion-independent); X'X is built one
  column at a time through the exact (count,k,n)@(count,n,1) path, with
  regression tests pinning the behavior
- planner: gpuAvailable detected from actual metallib presence (SwiftPM
  CLI cannot compile Metal shaders - swift test skips GPU suites, the
  xcodebuild-built app and tests get the GPU); reps >= 500 dispatches
  bootstrap to .gpu
- engine: backend-aware bootstrap, drops singular resamples, output
  titled 'GPU batched' with gpu scalar
- 68 tests green under xcodebuild (66 + GPU suites under swift test skip)

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

simon-pierre boucher committed 6 days ago (Aug 5, 2026) parent b24d834

Showing 7 changed files with +699 and −49

modified MetrikaKit/Package.resolved +19 −1

		`@@ -1,5 +1,5 @@`
1	1	{
2		− "originHash" : "e4ae512a67bec7ae66f2abd23f1e87240c39d28811734cdfafc7a91e4323f6db",
	2	+ "originHash" : "c4143a21eaf19f86159743269f92608859217bba0735131f48d3b12086f2b0c2",
3	3	"pins" : [
4	4	{
5	5	"identity" : "duckdb-swift",
		`@@ -10,6 +10,15 @@`
10	10	"version" : "1.1.3"
11	11	}
12	12	},
	13	+ {
	14	+ "identity" : "mlx-swift",
	15	+ "kind" : "remoteSourceControl",
	16	+ "location" : "https://github.com/ml-explore/mlx-swift",
	17	+ "state" : {
	18	+ "revision" : "0bb916c67f4b9e5c682cbe02a42c701c93ab5021",
	19	+ "version" : "0.31.6"
	20	+ }
	21	+ },
13	22	{
14	23	"identity" : "swift-argument-parser",
15	24	"kind" : "remoteSourceControl",
		`@@ -27,6 +36,15 @@`
27	36	"revision" : "a0cb0954ecb21e4e31b0070e6ed5674e8556685a",
28	37	"version" : "1.6.0"
29	38	}
	39	+ },
	40	+ {
	41	+ "identity" : "swift-numerics",
	42	+ "kind" : "remoteSourceControl",
	43	+ "location" : "https://github.com/apple/swift-numerics",
	44	+ "state" : {
	45	+ "revision" : "0c0290ff6b24942dadb83a929ffaaa1481df04a2",
	46	+ "version" : "1.1.1"
	47	+ }
30	48	}
31	49	],
32	50	"version" : 3

modified MetrikaKit/Package.swift +5 −1

		`@@ -31,6 +31,7 @@ let package = Package(`
31	31	.package(url: "https://github.com/duckdb/duckdb-swift", from: "1.0.0"),
32	32	.package(url: "https://github.com/apple/swift-collections", from: "1.1.0"),
33	33	.package(url: "https://github.com/apple/swift-argument-parser", from: "1.3.0"),
	34	+ .package(url: "https://github.com/ml-explore/mlx-swift", from: "0.18.0"),
34	35	],
35	36	targets: [
36	37	.target(
		`@@ -52,7 +53,10 @@ let package = Package(`
52	53	),
53	54	.target(
54	55	name: "ZQGPU",
55		− dependencies: ["ZQData"],
	56	+ dependencies: [
	57	+ "ZQData",
	58	+ .product(name: "MLX", package: "mlx-swift"),
	59	+ ],
56	60	swiftSettings: strictConcurrency
57	61	),
58	62	.target(

modified MetrikaKit/Sources/ZQEngine/Session.swift +81 −47

		`@@ -35,7 +35,7 @@ public actor ZQSession {`
35	35
36	36	private let store: ZQDataStore
37	37	private let parser = ZQCommandParser()
38		− private let planner = ZQPlanner()
	38	+ private let planner = ZQPlanner(gpuAvailable: ZQGPUBootstrap.isAvailable)
39	39	private let scriptCommands: [String: ZQScriptCommand]
40	40	private var logFileURL: URL?
41	41	private var scriptDepth = 0
		`@@ -98,10 +98,12 @@ public actor ZQSession {`
98	98	// User script commands shadow nothing built-in; check after parse
99	99	// failure would be better UX, but the verb table already resolved.
100	100	let plan = planner.plan(command, rowCount: frame.rowCount)
101		− return try await dispatch(plan.command)
	101	+ return try await dispatch(plan.command, backend: plan.backend)
102	102	}
103	103
104		− private func dispatch(_ command: ZQCommand) async throws -> ZQResult {
	104	+ private func dispatch(
	105	+ _ command: ZQCommand, backend: ZQBackend = .cpu
	106	+ ) async throws -> ZQResult {
105	107	if let script = scriptCommands[command.verb] {
106	108	guard scriptDepth < 8 else {
107	109	throw ZQEngineError("user command recursion too deep")
		`@@ -132,7 +134,7 @@ public actor ZQSession {`
132	134	case "set": return try handleSet(command)
133	135	case "xtset": return try handleXTSet(command)
134	136	case "display": return try handleDisplay(command)
135		− case "bootstrap": return try await handleBootstrap(command)
	137	+ case "bootstrap": return try await handleBootstrap(command, backend: backend)
136	138	case "graph": return try handleGraph(command)
137	139	case "log": return try handleLog(command)
138	140	default:
		`@@ -1010,7 +1012,9 @@ public actor ZQSession {`
1010	1012
1011	1013	// MARK: - Bootstrap prefix
1012	1014
1013		− private func handleBootstrap(_ command: ZQCommand) async throws -> ZQResult {
	1015	+ private func handleBootstrap(
	1016	+ _ command: ZQCommand, backend: ZQBackend
	1017	+ ) async throws -> ZQResult {
1014	1018	guard let body = command.body else {
1015	1019	throw ZQEngineError("bootstrap: syntax is 'bootstrap, reps(#): command'")
1016	1020	}
		`@@ -1045,57 +1049,83 @@ public actor ZQSession {`
1045	1049
1046	1050	// Pairs bootstrap: replicate r regenerates its indices from the
1047	1051	// Philox counter stream — deterministic and order-independent, so
1048		− // replicates can be computed in parallel chunks.
1049		− let chunkCount = min(
1050		− max(1, ProcessInfo.processInfo.activeProcessorCount), reps
1051		− )
1052		− let chunkSize = (reps + chunkCount - 1) / chunkCount
1053		− let draws: [[Double]] = try await withThrowingTaskGroup(
1054		− of: [(Int, [Double])].self
1055		− ) { group in
1056		− for chunk in 0..<chunkCount {
1057		− let lower = chunk * chunkSize
1058		− let upper = min(reps, lower + chunkSize)
1059		− guard lower < upper else { continue }
1060		− group.addTask {
1061		− var results: [(Int, [Double])] = []
1062		− results.reserveCapacity(upper - lower)
1063		− for replicate in lower..<upper {
1064		− let indices = ZQResampling.pairsBootstrapIndices(
1065		− replicate: replicate, sampleSize: n, generator: generator
1066		− )
1067		− let resampledY = indices.map { y[$0] }
1068		− let resampledPredictors = predictors.map { column in
1069		− (name: column.name, values: indices.map { column.values[$0] })
	1052	+ // replicates can run in parallel CPU chunks or as GPU batches with
	1053	+ // identical resamples.
	1054	+ let allDraws: [[Double]]
	1055	+ if backend == .gpu {
	1056	+ allDraws = ZQGPUBootstrap.pairsBootstrapOLS(
	1057	+ y: y,
	1058	+ predictors: predictors,
	1059	+ includeConstant: includeConstant,
	1060	+ replicates: reps,
	1061	+ seed: seed
	1062	+ )
	1063	+ } else {
	1064	+ allDraws = try await withThrowingTaskGroup(
	1065	+ of: [(Int, [Double])].self
	1066	+ ) { group in
	1067	+ let chunkCount = min(
	1068	+ max(1, ProcessInfo.processInfo.activeProcessorCount), reps
	1069	+ )
	1070	+ let chunkSize = (reps + chunkCount - 1) / chunkCount
	1071	+ for chunk in 0..<chunkCount {
	1072	+ let lower = chunk * chunkSize
	1073	+ let upper = min(reps, lower + chunkSize)
	1074	+ guard lower < upper else { continue }
	1075	+ group.addTask {
	1076	+ var results: [(Int, [Double])] = []
	1077	+ results.reserveCapacity(upper - lower)
	1078	+ for replicate in lower..<upper {
	1079	+ let indices = ZQResampling.pairsBootstrapIndices(
	1080	+ replicate: replicate, sampleSize: n, generator: generator
	1081	+ )
	1082	+ let resampledY = indices.map { y[$0] }
	1083	+ let resampledPredictors = predictors.map { column in
	1084	+ (name: column.name, values: indices.map { column.values[$0] })
	1085	+ }
	1086	+ let fit = try ZQOLS.fit(
	1087	+ y: resampledY,
	1088	+ predictors: resampledPredictors,
	1089	+ includeConstant: includeConstant,
	1090	+ variance: .classical
	1091	+ )
	1092	+ results.append((replicate, fit.coefficients.map(\.estimate)))
1070	1093	}
1071		− let fit = try ZQOLS.fit(
1072		− y: resampledY,
1073		− predictors: resampledPredictors,
1074		− includeConstant: includeConstant,
1075		− variance: .classical
1076		− )
1077		− results.append((replicate, fit.coefficients.map(\.estimate)))
	1094	+ return results
1078	1095	}
1079		− return results
1080	1096	}
1081		− }
1082		− var collected = [[Double]](repeating: [], count: reps)
1083		− for try await chunkResults in group {
1084		− for (replicate, estimates) in chunkResults {
1085		− collected[replicate] = estimates
	1097	+ var collected = [[Double]](repeating: [], count: reps)
	1098	+ for try await chunkResults in group {
	1099	+ for (replicate, estimates) in chunkResults {
	1100	+ collected[replicate] = estimates
	1101	+ }
1086	1102	}
	1103	+ return collected
1087	1104	}
1088		− return collected
	1105	+ }
	1106	+
	1107	+ // Singular resamples (NaN rows from the GPU path) are dropped,
	1108	+ // Stata-style.
	1109	+ let draws = allDraws.filter { row in !row.contains(where: \.isNaN) }
	1110	+ let completed = draws.count
	1111	+ guard completed >= 2 else {
	1112	+ throw ZQEngineError("bootstrap: too many failed replicates")
1089	1113	}
1090	1114
1091	1115	// Bootstrap standard errors: sd of replicate estimates.
	1116	+ let title = backend == .gpu
	1117	+ ? "Bootstrap results (pairs, GPU batched)"
	1118	+ : "Bootstrap results (pairs)"
1092	1119	var lines = [
1093		− "Bootstrap results (pairs) Replications = " +
1094		− TableFormatter.pad("\(reps)", 10),
	1120	+ TableFormatter.pad(title, 46, right: false) + "Replications = " +
	1121	+ TableFormatter.pad("\(completed)", 10),
1095	1122	" Number of obs = " +
1096	1123	TableFormatter.pad("\(n)", 10),
1097		− "",
1098	1124	]
	1125	+ if completed < reps {
	1126	+ lines.append("(\(reps - completed) replicates dropped: singular resample)")
	1127	+ }
	1128	+ lines.append("")
1099	1129	let widths = [12, 12, 12, 8, 8, 22]
1100	1130	lines.append(
1101	1131	TableFormatter.pad(sample.responseName, widths[0]) + " \| " +
		`@@ -1112,12 +1142,16 @@ public actor ZQSession {`
1112	1142	)
1113	1143	lines.append(TableFormatter.rule(widths))
1114	1144
1115		− var scalars: [String: Double] = ["reps": Double(reps), "N": Double(n)]
	1145	+ var scalars: [String: Double] = [
	1146	+ "reps": Double(completed),
	1147	+ "N": Double(n),
	1148	+ "gpu": backend == .gpu ? 1 : 0,
	1149	+ ]
1116	1150	for j in 0..<k {
1117	1151	let estimates = draws.map { $0[j] }
1118		− let mean = estimates.reduce(0, +) / Double(reps)
	1152	+ let mean = estimates.reduce(0, +) / Double(completed)
1119	1153	let variance = estimates.reduce(0) { $0 + ($1 - mean) * ($1 - mean) }
1120		− / Double(reps - 1)
	1154	+ / Double(completed - 1)
1121	1155	let se = variance.squareRoot()
1122	1156	let observed = point.coefficients[j].estimate
1123	1157	let z = observed / se

added MetrikaKit/Sources/ZQGPU/CholeskySolver.swift +62 −0

		`@@ -0,0 +1,62 @@`
	1	+//
	2	+// CholeskySolver.swift
	3	+// Metrika
	4	+//
	5	+// Author: Simon-Pierre Boucher
	6	+// Contact: contact@spboucher.ai
	7	+// Copyright © 2026 Simon-Pierre Boucher. All rights reserved.
	8	+//
	9	+
	10	+/// Plain-Swift Cholesky solve for the k×k normal equations of the GPU
	11	+/// bootstrap. k is the regressor count (single digits in practice), so a
	12	+/// dependency-free O(k³) routine beats reaching for LAPACK — and ZQStats
	13	+/// is the only module allowed to import Accelerate (CLAUDE.md §3).
	14	+enum CholeskySolver {
	15	+ struct SingularMatrix: Error {}
	16	+
	17	+ /// Solves A·x = b for symmetric positive-definite A (column-major,
	18	+ /// k×k). Throws on non-positive-definite input.
	19	+ static func solve(_ a: [Double], k: Int, rhs: [Double]) throws -> [Double] {
	20	+ precondition(a.count == k * k && rhs.count == k)
	21	+
	22	+ // Lower-triangular factor L with A = L·Lᵀ.
	23	+ var l = [Double](repeating: 0, count: k * k)
	24	+ for j in 0..<k {
	25	+ var diagonal = a[j * k + j]
	26	+ for p in 0..<j {
	27	+ diagonal -= l[p * k + j] * l[p * k + j]
	28	+ }
	29	+ guard diagonal > 0 else { throw SingularMatrix() }
	30	+ let root = diagonal.squareRoot()
	31	+ l[j * k + j] = root
	32	+ for i in (j + 1)..<k {
	33	+ var value = a[j * k + i]
	34	+ for p in 0..<j {
	35	+ value -= l[p * k + i] * l[p * k + j]
	36	+ }
	37	+ l[j * k + i] = value / root
	38	+ }
	39	+ }
	40	+
	41	+ // Forward substitution L·z = b.
	42	+ var z = [Double](repeating: 0, count: k)
	43	+ for i in 0..<k {
	44	+ var value = rhs[i]
	45	+ for p in 0..<i {
	46	+ value -= l[p * k + i] * z[p]
	47	+ }
	48	+ z[i] = value / l[i * k + i]
	49	+ }
	50	+
	51	+ // Back substitution Lᵀ·x = z.
	52	+ var x = [Double](repeating: 0, count: k)
	53	+ for i in stride(from: k - 1, through: 0, by: -1) {
	54	+ var value = z[i]
	55	+ for p in (i + 1)..<k {
	56	+ value -= l[i * k + p] * x[p]
	57	+ }
	58	+ x[i] = value / l[i * k + i]
	59	+ }
	60	+ return x
	61	+ }
	62	+}

added MetrikaKit/Sources/ZQGPU/MLXBootstrap.swift +340 −0

		`@@ -0,0 +1,340 @@`
	1	+//
	2	+// MLXBootstrap.swift
	3	+// Metrika
	4	+//
	5	+// Author: Simon-Pierre Boucher
	6	+// Contact: contact@spboucher.ai
	7	+// Copyright © 2026 Simon-Pierre Boucher. All rights reserved.
	8	+//
	9	+
	10	+import Foundation
	11	+import MLX
	12	+
	13	+/// GPU-batched pairs bootstrap (CLAUDE.md §5).
	14	+///
	15	+/// Resample indices come from Philox4x32-10 evaluated as vectorized MLX
	16	+/// ops with the same key schedule as the CPU reference, so GPU and CPU
	17	+/// streams are bit-identical for a given seed — the reproducibility
	18	+/// release blocker. The heavy O(reps·n·k²) work (gather + cross-products)
	19	+/// runs on the GPU in float32; the tiny O(reps·k³) normal-equation solves
	20	+/// run on the CPU in float64 (§5 hybrid).
	21	+///
	22	+/// Documented tolerance vs the CPU float64 path: coefficient draws agree
	23	+/// to ~1e-5 relative for well-conditioned designs (float32 gather sums).
	24	+public enum ZQGPUBootstrap {
	25	+
	26	+ /// Whether the Metal-backed MLX device is usable in this process.
	27	+ ///
	28	+ /// SwiftPM's command-line build cannot compile Metal shaders, so under
	29	+ /// `swift build`/`swift test` the mlx metallib does not exist and any
	30	+ /// MLX op would abort. Mirror the C++ loader's search (colocated,
	31	+ /// main-bundle Cmlx resource bundle, all loaded bundles) and report
	32	+ /// GPU availability only when a metallib is actually present —
	33	+ /// xcodebuild-built apps and tests have it, CLI builds fall back to
	34	+ /// the CPU path via the planner.
	35	+ public static let isAvailable: Bool = {
	36	+ #if arch(arm64)
	37	+ var roots: [URL] = []
	38	+ if let url = Bundle.main.resourceURL { roots.append(url) }
	39	+ roots.append(Bundle.main.bundleURL)
	40	+ if let executable = Bundle.main.executableURL {
	41	+ roots.append(executable.deletingLastPathComponent())
	42	+ }
	43	+ for bundle in Bundle.allBundles {
	44	+ if let url = bundle.resourceURL { roots.append(url) }
	45	+ }
	46	+ let candidates = roots.flatMap { root in
	47	+ [
	48	+ root.appendingPathComponent("mlx.metallib"),
	49	+ root.appendingPathComponent("default.metallib"),
	50	+ root.appendingPathComponent(
	51	+ "mlx-swift_Cmlx.bundle/Contents/Resources/default.metallib"
	52	+ ),
	53	+ ]
	54	+ }
	55	+ return candidates.contains { FileManager.default.fileExists(atPath: $0.path) }
	56	+ #else
	57	+ return false
	58	+ #endif
	59	+ }()
	60	+
	61	+ /// Runs `replicates` pairs-bootstrap OLS fits and returns the
	62	+ /// coefficient draws, one row per replicate (constant last when
	63	+ /// `includeConstant`). Replicates whose resampled design is singular
	64	+ /// come back as NaN rows for the caller to drop.
	65	+ public static func pairsBootstrapOLS(
	66	+ y: [Double],
	67	+ predictors: [(name: String, values: [Double])],
	68	+ includeConstant: Bool = true,
	69	+ replicates: Int,
	70	+ seed: UInt64,
	71	+ memoryBudgetBytes: Int = 512 << 20
	72	+ ) -> [[Double]] {
	73	+ let n = y.count
	74	+ let p = predictors.count
	75	+ let k = p + (includeConstant ? 1 : 0)
	76	+ precondition(n > k, "insufficient observations")
	77	+
	78	+ // Conditioning: the GPU stage forms X'X in float32, which squares
	79	+ // the condition number — raw data with an intercept loses 2–3
	80	+ // digits. Standardizing each predictor by its full-sample mean and
	81	+ // scale is an exact reparametrization (the shift/scale constants
	82	+ // are fixed across replicates), so the original coefficients are
	83	+ // recovered exactly after the float64 solve:
	84	+ // β_j = β̃_j / s_j, b₀ = α + ȳ − Σ c_j β_j.
	85	+ var shifts = [Double](repeating: 0, count: p)
	86	+ var scales = [Double](repeating: 1, count: p)
	87	+ var yShift = 0.0
	88	+ if includeConstant {
	89	+ yShift = y.reduce(0, +) / Double(n)
	90	+ for (j, column) in predictors.enumerated() {
	91	+ let mean = column.values.reduce(0, +) / Double(n)
	92	+ let variance = column.values.reduce(0) {
	93	+ $0 + ($1 - mean) * ($1 - mean)
	94	+ } / Double(n)
	95	+ shifts[j] = mean
	96	+ scales[j] = variance.squareRoot() > 1e-30 ? variance.squareRoot() : 1
	97	+ }
	98	+ }
	99	+
	100	+ // Standardized design matrix on the GPU, float32, row-major (n, k).
	101	+ var design = [Float]()
	102	+ design.reserveCapacity(n * k)
	103	+ for row in 0..<n {
	104	+ for (j, column) in predictors.enumerated() {
	105	+ design.append(Float((column.values[row] - shifts[j]) / scales[j]))
	106	+ }
	107	+ if includeConstant { design.append(1) }
	108	+ }
	109	+ let x = MLXArray(design, [n, k])
	110	+ let yGPU = MLXArray(y.map { Float($0 - yShift) }, [n, 1])
	111	+
	112	+ // §5 chunking: keep the gathered design under the memory budget.
	113	+ let bytesPerReplicate = n * (k + 1) * 4 * 2 // gathered X & y + transposes
	114	+ let chunkSize = max(1, min(replicates, memoryBudgetBytes / bytesPerReplicate))
	115	+
	116	+ var draws = [[Double]](repeating: [], count: replicates)
	117	+ var start = 0
	118	+ while start < replicates {
	119	+ let count = min(chunkSize, replicates - start)
	120	+ let (xtx, xty) = crossProducts(
	121	+ x: x, y: yGPU, n: n, k: k,
	122	+ firstReplicate: start, replicateCount: count, seed: seed
	123	+ )
	124	+ solveChunk(
	125	+ xtx: xtx, xty: xty, k: k,
	126	+ into: &draws, offset: start, count: count
	127	+ )
	128	+ start += count
	129	+ }
	130	+
	131	+ // Undo the standardization.
	132	+ if includeConstant {
	133	+ for r in 0..<replicates where !draws[r].isEmpty {
	134	+ var recovered = draws[r]
	135	+ var interceptShift = yShift
	136	+ for j in 0..<p {
	137	+ recovered[j] = draws[r][j] / scales[j]
	138	+ interceptShift -= shifts[j] * recovered[j]
	139	+ }
	140	+ recovered[p] = draws[r][p] + interceptShift
	141	+ draws[r] = recovered
	142	+ }
	143	+ }
	144	+ return draws
	145	+ }
	146	+
	147	+ /// One replicate's resample indices computed on the GPU — exposed so
	148	+ /// the cross-backend parity test can compare against the CPU
	149	+ /// reference (`ZQResampling.pairsBootstrapIndices`) bit-for-bit.
	150	+ public static func gpuIndices(
	151	+ replicate: Int, sampleSize: Int, seed: UInt64
	152	+ ) -> [Int] {
	153	+ let base = UInt64(replicate) * UInt64(sampleSize)
	154	+ let positions = MLXArray(0..<Int32(sampleSize)).asType(.uint64)
	155	+ + MLXArray(base, dtype: .uint64)
	156	+ let indices = philoxBoundedIndices(
	157	+ positions: positions, bound: sampleSize, seed: seed
	158	+ )
	159	+ eval(indices)
	160	+ return indices.asArray(Int32.self).map(Int.init)
	161	+ }
	162	+
	163	+ /// Diagnostic hook: one replicate's raw GPU cross-products with the
	164	+ /// same standardization as the production path.
	165	+ public static func debugCrossProducts(
	166	+ y: [Double], predictors: [(name: String, values: [Double])],
	167	+ replicate: Int, seed: UInt64
	168	+ ) -> (xtx: [Float], xty: [Float]) {
	169	+ let n = y.count
	170	+ let p = predictors.count
	171	+ let k = p + 1
	172	+ let meanY = y.reduce(0, +) / Double(n)
	173	+ var design = [Float]()
	174	+ for row in 0..<n {
	175	+ for column in predictors {
	176	+ let mean = column.values.reduce(0, +) / Double(n)
	177	+ let variance = column.values.reduce(0) {
	178	+ $0 + ($1 - mean) * ($1 - mean)
	179	+ } / Double(n)
	180	+ design.append(Float((column.values[row] - mean) / variance.squareRoot()))
	181	+ }
	182	+ design.append(1)
	183	+ }
	184	+ let x = MLXArray(design, [n, k])
	185	+ let yGPU = MLXArray(y.map { Float($0 - meanY) }, [n, 1])
	186	+ return crossProducts(
	187	+ x: x, y: yGPU, n: n, k: k,
	188	+ firstReplicate: replicate, replicateCount: 1, seed: seed
	189	+ )
	190	+ }
	191	+
	192	+ /// Diagnostic hook: cross-products for replicate 0 computed inside a
	193	+ /// batch of `batchSize`, to expose batch-dependent kernel numerics.
	194	+ public static func debugBatchedCrossProducts(
	195	+ y: [Double], predictors: [(name: String, values: [Double])],
	196	+ batchSize: Int, seed: UInt64
	197	+ ) -> (xtx: [Float], xty: [Float]) {
	198	+ let n = y.count
	199	+ let p = predictors.count
	200	+ let k = p + 1
	201	+ let meanY = y.reduce(0, +) / Double(n)
	202	+ var design = [Float]()
	203	+ for row in 0..<n {
	204	+ for column in predictors {
	205	+ let mean = column.values.reduce(0, +) / Double(n)
	206	+ let variance = column.values.reduce(0) {
	207	+ $0 + ($1 - mean) * ($1 - mean)
	208	+ } / Double(n)
	209	+ design.append(Float((column.values[row] - mean) / variance.squareRoot()))
	210	+ }
	211	+ design.append(1)
	212	+ }
	213	+ let x = MLXArray(design, [n, k])
	214	+ let yGPU = MLXArray(y.map { Float($0 - meanY) }, [n, 1])
	215	+ let (xtx, xty) = crossProducts(
	216	+ x: x, y: yGPU, n: n, k: k,
	217	+ firstReplicate: 0, replicateCount: batchSize, seed: seed
	218	+ )
	219	+ return (Array(xtx.prefix(k * k)), Array(xty.prefix(k)))
	220	+ }
	221	+
	222	+ // MARK: - GPU stage
	223	+
	224	+ /// Gathers the resampled design for a chunk and returns float32
	225	+ /// cross-products, pulled back to the CPU: X'X (count·k·k) and X'y
	226	+ /// (count·k).
	227	+ private static func crossProducts(
	228	+ x: MLXArray, y: MLXArray, n: Int, k: Int,
	229	+ firstReplicate: Int, replicateCount: Int, seed: UInt64
	230	+ ) -> (xtx: [Float], xty: [Float]) {
	231	+ // Draw positions: replicate r uses counters r·n ..< (r+1)·n,
	232	+ // matching ZQResampling.pairsBootstrapIndices.
	233	+ let base = UInt64(firstReplicate) * UInt64(n)
	234	+ let total = replicateCount * n
	235	+ let positions = MLXArray(0..<Int32(total)).asType(.uint64)
	236	+ + MLXArray(base, dtype: .uint64)
	237	+
	238	+ let indices = philoxBoundedIndices(positions: positions, bound: n, seed: seed)
	239	+ .reshaped([replicateCount, n])
	240	+
	241	+ // Gather rows: (count, n, k) and (count, n, 1).
	242	+ let xg = x.take(indices.flattened(), axis: 0)
	243	+ .reshaped([replicateCount, n, k])
	244	+ let yg = y.take(indices.flattened(), axis: 0)
	245	+ .reshaped([replicateCount, n, 1])
	246	+
	247	+ // Batched cross-products. X'X is computed one column at a time
	248	+ // through (count,k,n)@(count,n,1) products: MLX's batched GEMM
	249	+ // mis-accumulates small k×k outputs at batch ≥ 2 (~6e-4 relative,
	250	+ // verified against float64 sums, independent of strides/fusion),
	251	+ // while the single-column path is float32-exact. k is single-digit
	252	+ // so the extra dispatches are negligible next to the gather.
	253	+ let xt = xg.transposed(0, 2, 1) // (count, k, n)
	254	+ let xColumns = split(xg, parts: k, axis: 2) // k × (count, n, 1)
	255	+ let xtxColumns = xColumns.map { matmul(xt, $0) }
	256	+ let xtx = concatenated(xtxColumns, axis: 2) // (count, k, k)
	257	+ let xty = matmul(xt, yg) // (count, k, 1)
	258	+ eval(xtx, xty)
	259	+
	260	+ return (xtx.asArray(Float.self), xty.asArray(Float.self))
	261	+ }
	262	+
	263	+ /// Philox4x32-10 as vectorized MLX ops. Given absolute draw positions,
	264	+ /// returns bounded int32 indices in 0..<bound — bit-identical to
	265	+ /// `Philox4x32.integer(at:bound:)` on the CPU.
	266	+ static func philoxBoundedIndices(
	267	+ positions: MLXArray, bound: Int, seed: UInt64
	268	+ ) -> MLXArray {
	269	+ let mask32 = MLXArray(UInt64(0xFFFF_FFFF), dtype: .uint64)
	270	+ let block = positions >> 2
	271	+ let lane = (positions & MLXArray(UInt64(3), dtype: .uint64)).asType(.uint32)
	272	+
	273	+ var c0 = (block & mask32)
	274	+ var c1 = (block >> 32)
	275	+ var c2 = MLXArray.zeros(positions.shape, dtype: .uint64)
	276	+ var c3 = MLXArray.zeros(positions.shape, dtype: .uint64)
	277	+
	278	+ // Key schedule precomputed on the CPU — identical to the
	279	+ // reference: key_r = (seedLo + r·W0, seedHi + r·W1) mod 2³².
	280	+ let seedLo = UInt32(truncatingIfNeeded: seed)
	281	+ let seedHi = UInt32(truncatingIfNeeded: seed >> 32)
	282	+ let m0 = MLXArray(UInt64(0xD251_1F53), dtype: .uint64)
	283	+ let m1 = MLXArray(UInt64(0xCD9E_8D57), dtype: .uint64)
	284	+
	285	+ for round in 0..<10 {
	286	+ let k0 = seedLo &+ UInt32(round) &* 0x9E37_79B9
	287	+ let k1 = seedHi &+ UInt32(round) &* 0xBB67_AE85
	288	+ let product0 = m0 * c0 // c0 < 2³² so exact
	289	+ let product1 = m1 * c2
	290	+ let high0 = product0 >> 32
	291	+ let low0 = product0 & mask32
	292	+ let high1 = product1 >> 32
	293	+ let low1 = product1 & mask32
	294	+ let newC0 = (high1 ^ c1 ^ MLXArray(UInt64(k0), dtype: .uint64)) & mask32
	295	+ let newC2 = (high0 ^ c3 ^ MLXArray(UInt64(k1), dtype: .uint64)) & mask32
	296	+ c0 = newC0
	297	+ c1 = low1
	298	+ c2 = newC2
	299	+ c3 = low0
	300	+ }
	301	+
	302	+ // Lane select, then bounded multiply-shift like the reference.
	303	+ let word = which(
	304	+ lane .== MLXArray(UInt32(0), dtype: .uint32), c0,
	305	+ which(
	306	+ lane .== MLXArray(UInt32(1), dtype: .uint32), c1,
	307	+ which(lane .== MLXArray(UInt32(2), dtype: .uint32), c2, c3)
	308	+ )
	309	+ )
	310	+ let bounded = (word * MLXArray(UInt64(bound), dtype: .uint64)) >> 32
	311	+ return bounded.asType(.int32)
	312	+ }
	313	+
	314	+ // MARK: - CPU stage
	315	+
	316	+ /// Solves the chunk's normal equations in float64 via Cholesky.
	317	+ /// Singular systems yield NaN rows.
	318	+ private static func solveChunk(
	319	+ xtx: [Float], xty: [Float], k: Int,
	320	+ into draws: inout [[Double]], offset: Int, count: Int
	321	+ ) {
	322	+ for r in 0..<count {
	323	+ var a = [Double](repeating: 0, count: k * k)
	324	+ var b = [Double](repeating: 0, count: k)
	325	+ // MLX buffers are row-major; LAPACK wants column-major. X'X is
	326	+ // symmetric so the transpose is free; X'y is a vector.
	327	+ for i in 0..<k {
	328	+ b[i] = Double(xty[r * k + i])
	329	+ for j in 0..<k {
	330	+ a[j * k + i] = Double(xtx[r * k * k + i * k + j])
	331	+ }
	332	+ }
	333	+ if let solution = try? CholeskySolver.solve(a, k: k, rhs: b) {
	334	+ draws[offset + r] = solution
	335	+ } else {
	336	+ draws[offset + r] = [Double](repeating: .nan, count: k)
	337	+ }
	338	+ }
	339	+ }
	340	+}

added MetrikaKit/Tests/MetrikaKitTests/GPUBootstrapTests.swift +119 −0

		`@@ -0,0 +1,119 @@`
	1	+//
	2	+// GPUBootstrapTests.swift
	3	+// Metrika
	4	+//
	5	+// Author: Simon-Pierre Boucher
	6	+// Contact: contact@spboucher.ai
	7	+// Copyright © 2026 Simon-Pierre Boucher. All rights reserved.
	8	+//
	9	+
	10	+import Foundation
	11	+import Testing
	12	+import ZQGPU
	13	+import ZQParser
	14	+import ZQPlanner
	15	+import ZQStats
	16	+
	17	+/// GPU bootstrap cross-checks (CLAUDE.md §9): same seed ⇒ identical
	18	+/// resample indices across backends; estimates within the documented
	19	+/// float32 tolerance.
	20	+@Suite("GPU bootstrap", .enabled(if: ZQGPUBootstrap.isAvailable))
	21	+struct GPUBootstrapTests {
	22	+
	23	+ @Test("MLX Philox indices are bit-identical to the CPU reference")
	24	+ func philoxParity() {
	25	+ let seed: UInt64 = 42
	26	+ let n = 137
	27	+ let generator = Philox4x32(seed: seed)
	28	+
	29	+ for replicate in [0, 1, 7, 1000] {
	30	+ let cpu = ZQResampling.pairsBootstrapIndices(
	31	+ replicate: replicate, sampleSize: n, generator: generator
	32	+ )
	33	+ let gpu = ZQGPUBootstrap.gpuIndices(replicate: replicate, sampleSize: n, seed: seed)
	34	+ #expect(cpu == gpu, "replicate \(replicate)")
	35	+ }
	36	+ }
	37	+
	38	+ @Test("GPU coefficient draws match CPU within float32 tolerance")
	39	+ func drawParity() throws {
	40	+ // Small synthetic problem, well conditioned.
	41	+ let n = 200
	42	+ let generator = Philox4x32(seed: 7)
	43	+ let x = (0..<n).map { i in 5 + 10 * generator.uniform(at: UInt64(i)) }
	44	+ let y = (0..<n).map { i in
	45	+ 2 + 0.5 * x[i] + (generator.uniform(at: UInt64(n + i)) - 0.5)
	46	+ }
	47	+ let reps = 50
	48	+
	49	+ let gpuDraws = ZQGPUBootstrap.pairsBootstrapOLS(
	50	+ y: y, predictors: [("x", x)], replicates: reps, seed: 42
	51	+ )
	52	+
	53	+ let bootstrapGenerator = Philox4x32(seed: 42)
	54	+ for replicate in 0..<reps {
	55	+ let indices = ZQResampling.pairsBootstrapIndices(
	56	+ replicate: replicate, sampleSize: n, generator: bootstrapGenerator
	57	+ )
	58	+ let fit = try ZQOLS.fit(
	59	+ y: indices.map { y[$0] },
	60	+ predictors: [("x", indices.map { x[$0] })]
	61	+ )
	62	+ for (j, coefficient) in fit.coefficients.enumerated() {
	63	+ expectClose(
	64	+ gpuDraws[replicate][j], coefficient.estimate,
	65	+ rtol: 1e-4, "replicate \(replicate) b[\(coefficient.name)]"
	66	+ )
	67	+ }
	68	+ }
	69	+ }
	70	+
	71	+ @Test("GPU bootstrap is deterministic for a fixed seed")
	72	+ func determinism() {
	73	+ let n = 100
	74	+ let generator = Philox4x32(seed: 3)
	75	+ let x = (0..<n).map { i in generator.uniform(at: UInt64(i)) }
	76	+ let y = (0..<n).map { i in x[i] + generator.uniform(at: UInt64(n + i)) }
	77	+
	78	+ let first = ZQGPUBootstrap.pairsBootstrapOLS(
	79	+ y: y, predictors: [("x", x)], replicates: 20, seed: 99
	80	+ )
	81	+ let second = ZQGPUBootstrap.pairsBootstrapOLS(
	82	+ y: y, predictors: [("x", x)], replicates: 20, seed: 99
	83	+ )
	84	+ #expect(first == second)
	85	+ }
	86	+
	87	+ @Test("chunking does not change results")
	88	+ func chunkingInvariance() {
	89	+ let n = 80
	90	+ let generator = Philox4x32(seed: 5)
	91	+ let x = (0..<n).map { i in generator.uniform(at: UInt64(i)) }
	92	+ let y = (0..<n).map { i in 2 * x[i] + generator.uniform(at: UInt64(n + i)) }
	93	+
	94	+ let oneChunk = ZQGPUBootstrap.pairsBootstrapOLS(
	95	+ y: y, predictors: [("x", x)], replicates: 30, seed: 11
	96	+ )
	97	+ // Budget so small every chunk holds a single replicate.
	98	+ let manyChunks = ZQGPUBootstrap.pairsBootstrapOLS(
	99	+ y: y, predictors: [("x", x)], replicates: 30, seed: 11,
	100	+ memoryBudgetBytes: 1
	101	+ )
	102	+ #expect(oneChunk == manyChunks)
	103	+ }
	104	+
	105	+ @Test("planner sends large-reps bootstrap to the GPU")
	106	+ func plannerDispatch() throws {
	107	+ let planner = ZQPlanner(gpuAvailable: true)
	108	+ let parser = ZQCommandParser()
	109	+
	110	+ let large = try #require(try parser.parse("bootstrap, reps(10000): reg y x"))
	111	+ #expect(planner.plan(large, rowCount: 1000).backend == .gpu)
	112	+
	113	+ let small = try #require(try parser.parse("bootstrap, reps(100): reg y x"))
	114	+ #expect(planner.plan(small, rowCount: 1000).backend == .cpu)
	115	+
	116	+ let noGPU = ZQPlanner(gpuAvailable: false)
	117	+ #expect(noGPU.plan(large, rowCount: 1000).backend == .cpu)
	118	+ }
	119	+}

added MetrikaKit/Tests/MetrikaKitTests/GPUDiagnosticTests.swift +73 −0

		`@@ -0,0 +1,73 @@`
	1	+//
	2	+// GPUDiagnosticTests.swift
	3	+// Metrika
	4	+//
	5	+// Author: Simon-Pierre Boucher
	6	+// Contact: contact@spboucher.ai
	7	+// Copyright © 2026 Simon-Pierre Boucher. All rights reserved.
	8	+//
	9	+
	10	+import Foundation
	11	+import Testing
	12	+import ZQGPU
	13	+
	14	+/// Numerical regression tests for the GPU cross-product stage. These
	15	+/// pin down an MLX batched-GEMM issue: small k×k outputs mis-accumulate
	16	+/// (~6e-4 relative) for batch sizes ≥ 2, which the column-wise X'X
	17	+/// computation in `crossProducts` works around. If these start failing,
	18	+/// the workaround regressed or the kernel changed.
	19	+@Suite("GPU numerics regression", .enabled(if: ZQGPUBootstrap.isAvailable))
	20	+struct GPUDiagnosticTests {
	21	+ let n = 200
	22	+ let x: [Double]
	23	+ let y: [Double]
	24	+
	25	+ init() {
	26	+ let generator = Philox4x32(seed: 7)
	27	+ let count = 200
	28	+ let xs = (0..<count).map { i in 5 + 10 * generator.uniform(at: UInt64(i)) }
	29	+ x = xs
	30	+ y = (0..<count).map { i in
	31	+ 2 + 0.5 * xs[i] + (generator.uniform(at: UInt64(count + i)) - 0.5)
	32	+ }
	33	+ }
	34	+
	35	+ @Test("batched cross-products match float64 sums of identical data")
	36	+ func crossProductPrecision() {
	37	+ // Exact float64 sums of the float32-quantized data the GPU sees.
	38	+ let indices = ZQResampling.pairsBootstrapIndices(
	39	+ replicate: 0, sampleSize: n, generator: Philox4x32(seed: 42)
	40	+ )
	41	+ let meanX = x.reduce(0, +) / Double(n)
	42	+ let varX = x.reduce(0) { $0 + ($1 - meanX) * ($1 - meanX) } / Double(n)
	43	+ let meanY = y.reduce(0, +) / Double(n)
	44	+ let xs = indices.map { Float((x[$0] - meanX) / varX.squareRoot()) }
	45	+ let ys = indices.map { Float(y[$0] - meanY) }
	46	+ var sxx = 0.0, sxy = 0.0
	47	+ for i in 0..<n {
	48	+ sxx += Double(xs[i]) * Double(xs[i])
	49	+ sxy += Double(xs[i]) * Double(ys[i])
	50	+ }
	51	+
	52	+ for batch in [1, 2, 16] {
	53	+ let sums = ZQGPUBootstrap.debugBatchedCrossProducts(
	54	+ y: y, predictors: [("x", x)], batchSize: batch, seed: 42
	55	+ )
	56	+ expectClose(Double(sums.xtx[0]), sxx, rtol: 1e-5, "sxx batch=\(batch)")
	57	+ expectClose(Double(sums.xty[0]), sxy, rtol: 1e-5, "sxy batch=\(batch)")
	58	+ }
	59	+ }
	60	+
	61	+ @Test("replicate draws are independent of batch size")
	62	+ func batchSizeInvariance() {
	63	+ let reference = ZQGPUBootstrap.pairsBootstrapOLS(
	64	+ y: y, predictors: [("x", x)], replicates: 1, seed: 42
	65	+ )[0]
	66	+ for batch in [2, 10, 50] {
	67	+ let draws = ZQGPUBootstrap.pairsBootstrapOLS(
	68	+ y: y, predictors: [("x", x)], replicates: batch, seed: 42
	69	+ )
	70	+ #expect(draws[0] == reference, "batch=\(batch)")
	71	+ }
	72	+ }
	73	+}
	74