Switch SD3 and Flux2 noise generation to TorchRandomSource

swift/Sources/CoreAIDiffusionPipeline/Pipelines/Flux2Pipeline.swift

@@ -224,7 +224,7 @@ public struct Flux2Pipeline: DiffusionPipeline {
         // 4. Generate noise [1, inChannels, spatialSide, spatialSide]
         let latentShape = [1, inChannels, spatialSide, spatialSide]
         let latentCount = latentShape.reduce(1, *)
-        let noise = generateNoise(count: latentCount, seed: configuration.seed)
+        let noise = generateNoise(count: latentCount, seed: configuration.seed, sourceType: .torch)
         let noisePacked = packLatentsSpatialFlatten(
             noise, channels: inChannels, height: spatialSide, width: spatialSide)
 

swift/Sources/CoreAIDiffusionPipeline/Pipelines/SD3Pipeline.swift

@@ -109,7 +109,8 @@ public struct SD3Pipeline: DiffusionPipeline {
         // 3. Initial noise
         var latents = generateNoise(
             count: latentShape.reduce(1, *),
-            seed: configuration.seed)
+            seed: configuration.seed,
+            sourceType: .torch)
 
         // 4. Scheduler (SD3 flow matching; plain shift, no dynamic mu)
         let scheduler = DiscreteFlowScheduler(

swift/Sources/CoreAIDiffusionPipeline/RNG/NoiseGeneration.swift

@@ -11,14 +11,16 @@ public enum RandomSourceType: Sendable {
 }
 
 /// Generate Gaussian noise (mean 0, stdev 1) using the specified random source.
+/// `.torch` routes through `normalArray` to match `torch.randn` (batch-16 fill),
+/// not the scalar `nextNormal()` loop.
 public func generateNoise(count: Int, seed: UInt32, sourceType: RandomSourceType = .numPy) -> [Float] {
     switch sourceType {
     case .numPy:
         var rng = NumPyRandomSource(seed: seed)
         return (0..<count).map { _ in Float(rng.nextNormal()) }
     case .torch:
         var rng = TorchRandomSource(seed: seed)
-        return (0..<count).map { _ in Float(rng.nextNormal()) }
+        return rng.normalArray([count])
     case .nvidia:
         var rng = NvRandomSource(seed: seed)
         return (0..<count).map { _ in Float(rng.nextNormal()) }

swift/Sources/CoreAIDiffusionPipeline/RNG/TorchRandomSource.swift

@@ -92,38 +92,45 @@ public struct TorchRandomSource: RandomNumberGenerator, RandomSource, Sendable {
         nextGauss() * stdev + mean
     }
 
-    /// Matches torch.randn([shape], dtype=torch.float) behavior including
-    /// the batch-16 Box-Muller optimization for arrays >= 16 elements.
+    /// Matches `torch.randn(shape, dtype=.float)`, including PyTorch's `normal_fill_16`
+    /// batch-16 Box-Muller path for counts >= 16 (see the step comments below).
     public mutating func normalArray(_ shape: [Int], mean: Double = 0.0, stdev: Double = 1.0) -> [Float] {
         let count = shape.reduce(1, *)
         guard count >= 16 else {
             return (0..<count).map { _ in Float(nextNormal(mean: mean, stdev: stdev)) }
         }
+
+        // Step 1: Fill all elements with float32 uniforms (matches PyTorch's initial fill).
         var data = (0..<count).map { _ in Double(nextFloat()) }
-        for i in stride(from: 0, to: count - 15, by: 16) {
+
+        // Step 2: Box-Muller on complete 16-element blocks.
+        for blockStart in stride(from: 0, to: count - 15, by: 16) {
             for j in 0..<8 {
-                let u1 = 1 - data[i + j]
-                let u2 = data[i + j + 8]
+                let u1 = 1 - data[blockStart + j]
+                let u2 = data[blockStart + j + 8]
                 let radius = sqrt(-2.0 * log(u1))
                 let theta = 2.0 * .pi * u2
-                data[i + j] = radius * cos(theta) * stdev + mean
-                data[i + j + 8] = radius * sin(theta) * stdev + mean
+                data[blockStart + j] = radius * cos(theta) * stdev + mean
+                data[blockStart + j + 8] = radius * sin(theta) * stdev + mean
             }
         }
+
+        // Step 3: Remainder — re-fill last 16 with fresh uniforms, re-apply Box-Muller.
         if count % 16 != 0 {
-            for i in (count - 16)..<count {
-                data[i] = nextDouble()
+            let blockStart = count - 16
+            for i in 0..<16 {
+                data[blockStart + i] = Double(nextFloat())
             }
-            let i = count - 16
             for j in 0..<8 {
-                let u1 = 1 - data[i + j]
-                let u2 = data[i + j + 8]
+                let u1 = 1 - data[blockStart + j]
+                let u2 = data[blockStart + j + 8]
                 let radius = sqrt(-2.0 * log(u1))
                 let theta = 2.0 * .pi * u2
-                data[i + j] = radius * cos(theta) * stdev + mean
-                data[i + j + 8] = radius * sin(theta) * stdev + mean
+                data[blockStart + j] = radius * cos(theta) * stdev + mean
+                data[blockStart + j + 8] = radius * sin(theta) * stdev + mean
             }
         }
+
         return data.map { Float($0) }
     }
 }

swift/Tests/DiffusionPipelineTests/SchedulerTests.swift

@@ -75,6 +75,121 @@ struct RNGTests {
         let mean = samples.reduce(0, +) / Float(samples.count)
         #expect(abs(mean - 5.0) < 0.2)
     }
+
+    // MARK: - Torch RNG parity with Python torch.randn
+    //
+    // Expected values come from PyTorch's CPU generator. To regenerate an array, run its
+    // `// Reference:` snippet under `torch.Generator().manual_seed(<seed>)`.
+
+    // Reference: torch.manual_seed(42); [torch.randn(1, dtype=torch.float64).item() for _ in range(8)]
+    @Test("Torch scalar path matches Python reference (seed=42)")
+    func torchScalarParity() {
+        var rng = TorchRandomSource(seed: 42)
+        let expected: [Double] = [
+            0.3366903544, 0.1288094051, 0.2344623634, 0.2303330279,
+            -1.1228563767, -0.1863282993, 2.2082013356, -0.6379970568,
+        ]
+        for (i, exp) in expected.enumerated() {
+            let got = rng.nextNormal(mean: 0, stdev: 1)
+            #expect(abs(got - exp) < 1e-6, "scalar[\(i)]: got \(got), expected \(exp)")
+        }
+    }
+
+    // Reference: torch.manual_seed(42); torch.randn(32, dtype=torch.float32)
+    @Test("Torch batch path matches Python torch.randn (seed=42, count=32)")
+    func torchBatchParity32() {
+        var rng = TorchRandomSource(seed: 42)
+        let expected: [Float] = [
+            1.9269150496, 1.4872841835, 0.9007171988, -2.1055214405,
+            0.6784184575, -1.2345449924, -0.0430674814, -1.6046669483,
+            -0.7521361709, 1.6487228870, -0.3924786448, -1.4036067724,
+            -0.7278812528, -0.5594298840, -0.7688389421, 0.7624453902,
+            1.6423169374, -0.1595973223, -0.4973974824, 0.4395892322,
+            -0.7581311464, 1.0783176422, 0.8008005023, 1.6806205511,
+            1.2791243792, 1.2964228392, 0.6104664803, 1.3347377777,
+            -0.2316243201, 0.0417594910, -0.2515752614, 0.8598585129,
+        ]
+        let got = rng.normalArray([32], mean: 0, stdev: 1)
+        for (i, (g, e)) in zip(got, expected).enumerated() {
+            #expect(abs(g - e) < 1e-4, "batch[\(i)]: got \(g), expected \(e)")
+        }
+    }
+
+    // Reference: torch.manual_seed(42); torch.randn(16, dtype=torch.float32)
+    @Test("Torch batch path boundary (exactly 16 elements)")
+    func torchBatchBoundary16() {
+        var rng = TorchRandomSource(seed: 42)
+        let expected: [Float] = [
+            1.9269150496, 1.4872841835, 0.9007171988, -2.1055214405,
+            0.6784184575, -1.2345449924, -0.0430674814, -1.6046669483,
+            -0.7521361709, 1.6487228870, -0.3924786448, -1.4036067724,
+            -0.7278812528, -0.5594298840, -0.7688389421, 0.7624453902,
+        ]
+        let got = rng.normalArray([16], mean: 0, stdev: 1)
+        for (i, (g, e)) in zip(got, expected).enumerated() {
+            #expect(abs(g - e) < 1e-4, "boundary[\(i)]: got \(g), expected \(e)")
+        }
+    }
+
+    // Reference: torch.manual_seed(42); torch.randn(17, dtype=torch.float32)
+    @Test("Torch batch path remainder (17 elements)")
+    func torchBatchRemainder17() {
+        var rng = TorchRandomSource(seed: 42)
+        let expected: [Float] = [
+            1.9269150496, -0.1595973223, -0.4973974824, 0.4395892322,
+            -0.7581311464, 1.0783176422, 0.8008005023, 1.6806205511,
+            0.3558597863, 1.2964228392, 0.6104664803, 1.3347377777,
+            -0.2316243201, 0.0417594910, -0.2515752614, 0.8598585129,
+            -0.3097269237,
+        ]
+        let got = rng.normalArray([17], mean: 0, stdev: 1)
+        for (i, (g, e)) in zip(got, expected).enumerated() {
+            #expect(abs(g - e) < 1e-4, "remainder[\(i)]: got \(g), expected \(e)")
+        }
+    }
+
+    // Reference: torch.manual_seed(0); torch.randn(32, dtype=torch.float32)
+    @Test("Torch batch path matches Python torch.randn (seed=0, count=32)")
+    func torchBatchParity0() {
+        var rng = TorchRandomSource(seed: 0)
+        let expected: [Float] = [
+            -1.1258398294, -1.1523602009, -0.2505785823, -0.4338788390,
+            0.8487103581, 0.6920092106, -0.3160127699, -2.1152195930,
+            0.3222749233, -1.2633347511, 0.3499831855, 0.3081339002,
+            0.1198415086, 1.2376579046, 1.1167771816, -0.2472776473,
+            -1.3526537418, -1.6959313154, 0.5666505098, 0.7935084105,
+            0.5988394618, -1.5550950766, -0.3413603008, 1.8530061245,
+            0.7501894236, -0.5854971409, -0.1733970195, 0.1834779233,
+            1.3893661499, 1.5863343477, 0.9462983608, -0.8436768055,
+        ]
+        let got = rng.normalArray([32], mean: 0, stdev: 1)
+        for (i, (g, e)) in zip(got, expected).enumerated() {
+            #expect(abs(g - e) < 1e-4, "batch[\(i)]: got \(g), expected \(e)")
+        }
+    }
+
+    // Reference: torch.manual_seed(42); t = torch.randn([1,16,16,16], dtype=torch.float32)
+    @Test("Torch batch path realistic shape (4096 elements, seed=42)")
+    func torchBatchRealisticShape() {
+        var rng = TorchRandomSource(seed: 42)
+        let got = rng.normalArray([1, 16, 16, 16], mean: 0, stdev: 1)
+        #expect(got.count == 4096)
+
+        let expectedFirst: [Float] = [
+            1.9269150496, 1.4872841835, 0.9007171988, -2.1055214405,
+            0.6784184575, -1.2345449924, -0.0430674814, -1.6046669483,
+        ]
+        let expectedLast: [Float] = [
+            1.5869791508, 0.1421326697, 0.3760589659, -0.7916260362,
+            2.6677629948, -0.1403129250, 0.9416193962, -0.0118428767,
+        ]
+        for (i, (g, e)) in zip(got.prefix(8), expectedFirst).enumerated() {
+            #expect(abs(g - e) < 1e-4, "first[\(i)]: got \(g), expected \(e)")
+        }
+        for (i, (g, e)) in zip(got.suffix(8), expectedLast).enumerated() {
+            #expect(abs(g - e) < 1e-4, "last[\(i)]: got \(g), expected \(e)")
+        }
+    }
 }
 
 @Suite("Schedulers")