Enhanced the random seed generation function. (#328)

Author: eaplatanios

Sources/TensorFlow/Core/Utilities.swift

@@ -167,3 +167,152 @@ extension UnsafeMutablePointer where Pointee == CTensorHandle? {
         self.init(unwrapped)
     }
 }
+
+//===------------------------------------------------------------------------------------------===//
+// Hashing
+//===------------------------------------------------------------------------------------------===//
+
+internal extension FixedWidthInteger {
+    init(bytes: ArraySlice<UInt8>, startingAt index: Int) {
+        if bytes.isEmpty { self.init(0); return }
+        let count = bytes.count
+        self.init(0)
+        for i in 0..<MemoryLayout<Self>.size {
+            let j = (MemoryLayout<Self>.size - i - 1) * 8
+            self |= count > 0 ? Self(bytes[index.advanced(by: i)]) << j : 0
+        }
+    }
+
+    func bytes(count byteCount: Int = MemoryLayout<Self>.size) -> [UInt8] {
+        var littleEndianValue = littleEndian
+        return withUnsafePointer(to: &littleEndianValue) { pointer -> [UInt8] in
+            let bytesPointer = UnsafeMutablePointer<UInt8>(OpaquePointer(pointer))
+            var bytes = [UInt8](repeating: 0, count: byteCount)
+            for i in 0..<Swift.min(MemoryLayout<Self>.size, byteCount) {
+                bytes[byteCount - 1 - i] = (bytesPointer + i).pointee
+            }
+            return bytes
+        }
+    }
+}
+
+internal extension Array where Element == UInt8 {
+    func sha512() -> SIMD64<UInt8> {
+        // First we define some useful constants.
+        let blockSize = 128
+        let k: [UInt64] = [
+            0x428a2f98d728ae22, 0x7137449123ef65cd, 0xb5c0fbcfec4d3b2f, 0xe9b5dba58189dbbc,
+            0x3956c25bf348b538, 0x59f111f1b605d019, 0x923f82a4af194f9b, 0xab1c5ed5da6d8118,
+            0xd807aa98a3030242, 0x12835b0145706fbe, 0x243185be4ee4b28c, 0x550c7dc3d5ffb4e2,
+            0x72be5d74f27b896f, 0x80deb1fe3b1696b1, 0x9bdc06a725c71235, 0xc19bf174cf692694,
+            0xe49b69c19ef14ad2, 0xefbe4786384f25e3, 0x0fc19dc68b8cd5b5, 0x240ca1cc77ac9c65,
+            0x2de92c6f592b0275, 0x4a7484aa6ea6e483, 0x5cb0a9dcbd41fbd4, 0x76f988da831153b5,
+            0x983e5152ee66dfab, 0xa831c66d2db43210, 0xb00327c898fb213f, 0xbf597fc7beef0ee4,
+            0xc6e00bf33da88fc2, 0xd5a79147930aa725, 0x06ca6351e003826f, 0x142929670a0e6e70,
+            0x27b70a8546d22ffc, 0x2e1b21385c26c926, 0x4d2c6dfc5ac42aed, 0x53380d139d95b3df,
+            0x650a73548baf63de, 0x766a0abb3c77b2a8, 0x81c2c92e47edaee6, 0x92722c851482353b,
+            0xa2bfe8a14cf10364, 0xa81a664bbc423001, 0xc24b8b70d0f89791, 0xc76c51a30654be30,
+            0xd192e819d6ef5218, 0xd69906245565a910, 0xf40e35855771202a, 0x106aa07032bbd1b8,
+            0x19a4c116b8d2d0c8, 0x1e376c085141ab53, 0x2748774cdf8eeb99, 0x34b0bcb5e19b48a8,
+            0x391c0cb3c5c95a63, 0x4ed8aa4ae3418acb, 0x5b9cca4f7763e373, 0x682e6ff3d6b2b8a3,
+            0x748f82ee5defb2fc, 0x78a5636f43172f60, 0x84c87814a1f0ab72, 0x8cc702081a6439ec,
+            0x90befffa23631e28, 0xa4506cebde82bde9, 0xbef9a3f7b2c67915, 0xc67178f2e372532b,
+            0xca273eceea26619c, 0xd186b8c721c0c207, 0xeada7dd6cde0eb1e, 0xf57d4f7fee6ed178,
+            0x06f067aa72176fba, 0x0a637dc5a2c898a6, 0x113f9804bef90dae, 0x1b710b35131c471b,
+            0x28db77f523047d84, 0x32caab7b40c72493, 0x3c9ebe0a15c9bebc, 0x431d67c49c100d4c,
+            0x4cc5d4becb3e42b6, 0x597f299cfc657e2a, 0x5fcb6fab3ad6faec, 0x6c44198c4a475817]
+
+        var accumulated = self
+        let lengthInBits = accumulated.count * 8
+        let lengthBytes = lengthInBits.bytes(count: blockSize / 8)
+        
+        // Step 1: Append padding.
+        let msgLength = accumulated.count
+        // Append one bit (`UInt8` with one bit) to the message.
+        accumulated.append(0x80)
+        // Append `0` bits until the length of `accumulated` in bits is 448 (mod 512).
+        let max = blockSize * 7 / 8
+        accumulated += [UInt8](
+            repeating: 0,
+            count: msgLength % blockSize < max ?
+                max - 1 - (msgLength % blockSize) :
+                blockSize + max - 1 - (msgLength % blockSize))
+
+        // Step 2: Append the message length as a 64-bit representation of `lengthInBits`.
+        accumulated += lengthBytes
+
+        // Step 3: Process the array bytes.
+        var accumulatedHash = SIMD8<UInt64>(
+            0x6a09e667f3bcc908, 0xbb67ae8584caa73b, 0x3c6ef372fe94f82b, 0xa54ff53a5f1d36f1, 
+            0x510e527fade682d1, 0x9b05688c2b3e6c1f, 0x1f83d9abfb41bd6b, 0x5be0cd19137e2179)
+        var index = 0
+        while index < accumulated.count {
+            let chunk = accumulated[index..<(index + blockSize)]
+            index += blockSize
+
+            // Break chunk into sixteen 64-bit words w[j], 0 ≤ j ≤ 15, in big-endian format.
+            // Extend the sixteen 64-bit words into eighty 64-bit words:
+            var w = [UInt64](repeating: 0, count: k.count)
+            for x in k.indices {
+                switch x {
+                case 0...15:
+                    let start = chunk.startIndex.advanced(by: x * 8)
+                    w[x] = UInt64(bytes: chunk, startingAt: start)
+                    break
+                default:
+                    let s0Term0 = ((w[x - 15] >> 1 ^ w[x - 15]) >> 6 ^ w[x - 15]) >> 1
+                    let s0Term1 = (w[x - 15] << 7 ^ w[x - 15]) << 56
+                    let s0 = s0Term0 ^ s0Term1
+                    let s1Term0 = ((w[x - 2] >> 42 ^ w[x - 2]) >> 13 ^ w[x - 2]) >> 6
+                    let s1Term1 = (w[x - 2] << 42 ^ w[x - 2]) << 3
+                    let s1 = s1Term0 ^ s1Term1
+                    w[x] = w[x - 16] &+ s0 &+ w[x - 7] &+ s1
+                    break
+                }
+            }
+
+            var hashCopy = accumulatedHash
+            for j in k.indices {
+                let s0Term0 = ((hashCopy[0] >> 5 ^ hashCopy[0]) >> 6 ^ hashCopy[0]) >> 28
+                let s0Term1 = ((hashCopy[0] << 6 ^ hashCopy[0]) << 5 ^ hashCopy[0]) << 25
+                let s0 = s0Term0 ^ s0Term1
+                let s1Term0 = ((hashCopy[4] >> 23 ^ hashCopy[4]) >> 4 ^ hashCopy[4]) >> 14
+                let s1Term1 = ((hashCopy[4] << 4 ^ hashCopy[4]) << 23 ^ hashCopy[4]) << 23
+                let s1 = s1Term0 ^ s1Term1
+                let maj = (hashCopy[0] & hashCopy[1]) ^ 
+                    (hashCopy[0] & hashCopy[2]) ^
+                    (hashCopy[1] & hashCopy[2])
+                let t2 = s0 &+ maj
+                let ch = (hashCopy[4] & hashCopy[5]) ^ (~hashCopy[4] & hashCopy[6])
+                let t1 = hashCopy[7] &+ s1 &+ ch &+ k[j] &+ w[j]
+                hashCopy[7] = hashCopy[6]
+                hashCopy[6] = hashCopy[5]
+                hashCopy[5] = hashCopy[4]
+                hashCopy[4] = hashCopy[3] &+ t1
+                hashCopy[3] = hashCopy[2]
+                hashCopy[2] = hashCopy[1]
+                hashCopy[1] = hashCopy[0]
+                hashCopy[0] = t1 &+ t2
+            }
+            accumulatedHash &+= hashCopy
+        }
+
+        // Step 4: Return the computed hash.
+        var result = SIMD64<UInt8>()
+        var position = 0
+        for index in accumulatedHash.indices {
+            let h = accumulatedHash[index]
+            result[position + 0] = UInt8((h >> 56) & 0xff)
+            result[position + 1] = UInt8((h >> 48) & 0xff)
+            result[position + 2] = UInt8((h >> 40) & 0xff)
+            result[position + 3] = UInt8((h >> 32) & 0xff)
+            result[position + 4] = UInt8((h >> 24) & 0xff)
+            result[position + 5] = UInt8((h >> 16) & 0xff)
+            result[position + 6] = UInt8((h >> 8) & 0xff)
+            result[position + 7] = UInt8(h & 0xff)
+            position += 8
+        }
+
+        return result
+    }
+}

Sources/TensorFlow/Random.swift

@@ -19,8 +19,31 @@ import Glibc
 #endif
 
 /// Generates a new random seed for TensorFlow.
-public func randomSeedForTensorFlow() -> (Int32, Int32) {
-    (Int32.random(in: Int32.min..<Int32.max), Int32.random(in: Int32.min..<Int32.max))
+public func randomSeedForTensorFlow(using seed: (Int32, Int32)? = nil) -> (Int32, Int32) {
+    var strongSeed = UInt64(0)
+    if let s = seed {
+        let bytes = (s.0.bytes() + s.1.bytes())[...]
+        let singleSeed = UInt64(bytes: bytes, startingAt: bytes.startIndex)
+        strongSeed = UInt64(pow(Double(singleSeed % 2), Double(8 * 8)))
+    } else {
+        strongSeed = UInt64.random(in: UInt64.min..<UInt64.max)
+    }
+
+    // Many machine learning systems are likely to have many random number generators active at
+    // once (e.g., in reinforcement learning we may have an environment running in multiple 
+    // processes). There is literature indicating that having linear correlations between seeds of 
+    // multiple PRNG's can correlate the outputs:
+    //   - http://blogs.unity3d.com/2015/01/07/a-primer-on-repeatable-random-numbers
+    //   - http://stackoverflow.com/questions/1554958/how-different-do-random-seeds-need-to-be
+    //   - http://dl.acm.org/citation.cfm?id=1276928
+    // Thus, for sanity we hash the generated seed before using it, This scheme is likely not 
+    // crypto-strength, but it should be good enough to get rid of simple correlations.
+    // Reference: https://github.com/openai/gym/blob/master/gym/utils/seeding.py
+
+    let hash = strongSeed.bytes().sha512()
+    let first = Int32(bytes: [hash[0], hash[1], hash[2], hash[3]], startingAt: 0)
+    let second = Int32(bytes: [hash[4], hash[5], hash[6], hash[7]], startingAt: 0)
+    return (first, second)
 }
 
 //===------------------------------------------------------------------------------------------===//

Tests/TensorFlowTests/CoreTests/UtilitiesTests.swift

@@ -0,0 +1,30 @@
+// Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+import XCTest
+@testable import TensorFlow
+
+final class UtilitiesTests: XCTestCase {
+    func testSHA512() {
+        XCTAssertEqual(
+            [UInt8](repeating: 0x61, count: 1000).sha512(),
+            SIMD64<UInt8>([
+                103, 186,  85,  53, 164, 110,  63, 134, 219, 251, 237, 140, 187, 175,   1,  37,
+                199, 110, 213,  73, 255, 139,  11, 158,   3, 224, 200, 140, 249,  15, 166,  52,
+                250, 123,  18, 180, 125, 119, 182, 148, 222,  72, 138, 206, 141, 154, 101, 150,
+                125, 201, 109, 245, 153, 114, 125,  50, 146, 168, 217, 212,  71, 112, 156, 151]))
+    }
+
+    static var allTests = [("testSHA512", testSHA512)]
+}

Tests/TensorFlowTests/XCTestManifests.swift

@@ -17,6 +17,7 @@ import XCTest
 #if !os(macOS)
 public func allTests() -> [XCTestCaseEntry] {
     return [
+        testCase(UtilitiesTests.allTests),
         testCase(LossTests.allTests),
         testCase(PRNGTests.allTests),
         testCase(TrivialModelTests.allTests),