spb/metrika Public

1//2//  GPUDiagnosticTests.swift3//  Metrika4//5//  Author:  Simon-Pierre Boucher6//  Contact: contact@spboucher.ai7//  Copyright © 2026 Simon-Pierre Boucher. All rights reserved.8//910import Foundation11import Testing12import ZQGPU1314/// Numerical regression tests for the GPU cross-product stage. These15/// pin down an MLX batched-GEMM issue: small k×k outputs mis-accumulate16/// (~6e-4 relative) for batch sizes ≥ 2, which the column-wise X'X17/// 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))20struct GPUDiagnosticTests {21    let n = 20022    let x: [Double]23    let y: [Double]2425    init() {26        let generator = Philox4x32(seed: 7)27        let count = 20028        let xs = (0..<count).map { i in 5 + 10 * generator.uniform(at: UInt64(i)) }29        x = xs30        y = (0..<count).map { i in31            2 + 0.5 * xs[i] + (generator.uniform(at: UInt64(count + i)) - 0.5)32        }33    }3435    @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.047        for i in 0..<n {48            sxx += Double(xs[i]) * Double(xs[i])49            sxy += Double(xs[i]) * Double(ys[i])50        }5152        for batch in [1, 2, 16] {53            let sums = ZQGPUBootstrap.debugBatchedCrossProducts(54                y: y, predictors: [("x", x)], batchSize: batch, seed: 4255            )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    }6061    @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: 4265        )[0]66        for batch in [2, 10, 50] {67            let draws = ZQGPUBootstrap.pairsBootstrapOLS(68                y: y, predictors: [("x", x)], replicates: batch, seed: 4269            )70            #expect(draws[0] == reference, "batch=\(batch)")71        }72    }73}74