wehub-resource-sync
92 lines
3.3 KiB
Swift
92 lines
3.3 KiB
Swift
import Foundation
|
|
import CoreML
|
|
|
|
/// A random source consistent with NVIDIA curandom
|
|
///
|
|
/// This implementation references to:
|
|
/// https://github.com/dsnz/random/blob/master/philox.py for Philox_M4_32 configuration.
|
|
///
|
|
@available(iOS 16.2, macOS 13.1, *)
|
|
struct NvRandomSource: RandomSource {
|
|
public let seed: UInt64
|
|
private var offset: UInt32
|
|
|
|
/// Initialize with a random seed
|
|
///
|
|
/// - Parameters
|
|
/// - seed: Seed for underlying Philox M4 32 generator
|
|
/// - Returns random source
|
|
public init(seed: UInt32) {
|
|
self.seed = UInt64(seed)
|
|
offset = 0
|
|
}
|
|
|
|
static private let PHILOX_M4_32: (UInt32, UInt32) = (0xD251_1F53, 0xCD9E_8D57)
|
|
static private let PHILOX_W_32: (UInt32, UInt32) = (0x9E37_79B9, 0xBB67_AE85)
|
|
|
|
static private func philox4Round(counter: inout [[UInt32]], key: [[UInt32]]) {
|
|
for i in 0..<counter[0].count {
|
|
let v1: UInt64 = UInt64(counter[0][i]) * UInt64(PHILOX_M4_32.0)
|
|
let v2: UInt64 = UInt64(counter[2][i]) * UInt64(PHILOX_M4_32.1)
|
|
counter[0][i] = UInt32(v2 >> 32) ^ counter[1][i] ^ key[0][i]
|
|
counter[1][i] = UInt32(v2 & 0xffff_ffff)
|
|
counter[2][i] = UInt32(v1 >> 32) ^ counter[3][i] ^ key[1][i]
|
|
counter[3][i] = UInt32(v1 & 0xffff_ffff)
|
|
}
|
|
}
|
|
|
|
static private func philox4Bumpkey(key: inout [[UInt32]]) {
|
|
for (i, element) in key[0].enumerated() {
|
|
key[0][i] = element &+ PHILOX_W_32.0
|
|
}
|
|
for (i, element) in key[1].enumerated() {
|
|
key[1][i] = element &+ PHILOX_W_32.1
|
|
}
|
|
}
|
|
|
|
static private func philox4_32(counter: inout [[UInt32]], key: inout [[UInt32]], rounds: Int = 10) {
|
|
for _ in 0..<(rounds - 1) {
|
|
philox4Round(counter: &counter, key: key)
|
|
philox4Bumpkey(key: &key)
|
|
}
|
|
philox4Round(counter: &counter, key: key)
|
|
}
|
|
|
|
private func boxMuller(_ counter1: [UInt32], _ counter2: [UInt32], mean: Double, stdev: Double) -> [Double] {
|
|
// Box-Muller transform
|
|
return zip(counter1, counter2).map {
|
|
let u: Double = Double($0) / 4294967296.0 + (1.0 / 8589934592.0)
|
|
let v: Double = Double($1) * (.pi / 2147483648.0) + (.pi / 4294967296.0)
|
|
let radius = stdev * sqrt(-2.0 * log(u))
|
|
return radius * sin(v) + mean
|
|
}
|
|
}
|
|
|
|
private mutating func normalArray(count: Int, mean: Double, stdev: Double) -> [Double] {
|
|
var counter: [[UInt32]] = [
|
|
Array(repeating: offset, count: count),
|
|
Array(repeating: 0, count: count),
|
|
Array(0..<UInt32(count)),
|
|
Array(repeating: 0, count: count),
|
|
]
|
|
offset += 1
|
|
var key: [[UInt32]] = [
|
|
Array(repeating: UInt32(seed & 0xffff_ffff), count: count),
|
|
Array(repeating: UInt32(seed >> 32), count: count),
|
|
]
|
|
Self.philox4_32(counter: &counter, key: &key)
|
|
return boxMuller(counter[0], counter[1], mean: mean, stdev: stdev)
|
|
}
|
|
|
|
/// Generates a random value from a normal distribution with given mean and standard deviation.
|
|
mutating func nextNormal(mean: Double = 0.0, stdev: Double = 1.0) -> Double {
|
|
return normalArray(count: 1, mean: mean, stdev: stdev)[0]
|
|
}
|
|
|
|
/// Generate a shaped array with scalars from a normal distribution with given mean and standard deviation.
|
|
mutating func normalShapedArray(_ shape: [Int], mean: Double = 0.0, stdev: Double = 1.0) -> MLShapedArray<Double> {
|
|
let count = shape.reduce(1, *)
|
|
return .init(scalars: normalArray(count: count, mean: mean, stdev: stdev), shape: shape)
|
|
}
|
|
}
|