Add implementation for RNG-MRG, using MRG31k3p random number generator.

myia/rng_mrg.py

@@ -0,0 +1,290 @@
+"""
+Implementation of MRG31k3p random number generator for Myia.
+
+Based on Theano implementation (2019/11/05):
+https://github.com/Theano/Theano/blob/master/theano/sandbox/rng_mrg.py
+
+While Theano implementation uses many parallel streams in a matrix to
+generate numbers, current myia implementation below uses only
+1 stream vector for simplicity.
+
+Generator code in SSJ package (L'Ecuyer & Simard):
+http://www.iro.umontreal.ca/~simardr/ssj/indexe.html
+Page up-to-date (2019/11/05):
+http://simul.iro.umontreal.ca/ssj/indexe.html
+Original Github project page (2019/11/05):
+https://github.com/umontreal-simul/ssj
+Original JAVA implementation on Github (2019/11/05):
+https://github.com/umontreal-simul/ssj/blob/master/src/main/java/umontreal/ssj/rng/MRG31k3p.java
+
+The MRG31k3p algorithm was published in:
+
+P. L'Ecuyer and R. Touzin, Fast Combined Multiple Recursive Generators with
+Multipliers of the form a = +/- 2^d +/- 2^e, Proceedings of the 2000 Winter
+Simulation Conference, Dec. 2000, 683-689.
+
+The conception of the multi-stream from MRG31k3p was published in:
+
+P. L'Ecuyer and R. Simard and E. Jack Chen and W. David Kelton,
+An Object-Oriented Random-Number Package with Many Long Streams
+and Substreams, Operations Research, volume 50, number 6, 2002, 1073-1075.
+"""
+
+from dataclasses import dataclass
+
+import numpy as np
+
+# ----------
+# Constants.
+# ----------
+
+M1 = np.int32(2147483647)  # 2^31 - 1
+M2 = np.int32(2147462579)  # 2^31 - 21069
+MASK12 = np.int32(511)  # 2^9 - 1
+MASK13 = np.int32(16777215)  # 2^24 - 1
+MASK2 = np.int32(65535)  # 2^16 - 1
+MULT2 = np.int32(21069)
+A1p72 = np.asarray([[1516919229, 758510237, 499121365],
+                    [1884998244, 1516919229, 335398200],
+                    [601897748, 1884998244, 358115744]],
+                   dtype='int64')
+A2p72 = np.asarray([[1228857673, 1496414766, 954677935],
+                    [1133297478, 1407477216, 1496414766],
+                    [2002613992, 1639496704, 1407477216]],
+                   dtype='int64')
+A1p134 = np.asarray([[1702500920, 1849582496, 1656874625],
+                     [828554832, 1702500920, 1512419905],
+                     [1143731069, 828554832, 102237247]],
+                    dtype='int64')
+A2p134 = np.asarray([[796789021, 1464208080, 607337906],
+                     [1241679051, 1431130166, 1464208080],
+                     [1401213391, 1178684362, 1431130166]],
+                    dtype='int64')
+
+i0 = np.int32(0)
+i7 = np.int32(7)
+i9 = np.int32(9)
+i15 = np.int32(15)
+i16 = np.int32(16)
+i22 = np.int32(22)
+i24 = np.int32(24)
+
+# Hashable versions of useful nd-arrays above
+# (to avoid hash error during myia compilation).
+a_1_p_134 = tuple(
+    tuple(np.int64(A1p134[i, j]) for j in range(3)) for i in range(3))
+a_2_p_134 = tuple(
+    tuple(np.int64(A2p134[i, j]) for j in range(3)) for i in range(3))
+
+
+# Tuple of d-types supported by RNG, and index for each type in tuple.
+SUPP_TYPES = (np.float16, np.float32, np.float64)
+FLOAT16, FLOAT32, FLOAT64 = 0, 1, 2
+
+
+# ------------------
+# Private functions.
+# ------------------
+
+def _mat_vec_mod_m(a, i, j, k, m):
+    """
+    Specific matrix-vector operation used to increment random state.
+
+    Given 3x3 matrix `a` and vector v = (i, j, k), compute:
+        output = numpy.sum((a * v) % m, axis=1) % m
+    :param a: a 2-d tensor (matrix) with shape (3, 3)
+    :param i: a scalar
+    :param j: a scalar
+    :param k: a scalar
+    :param m: a scalar to use for modulo
+    :return: a vector with 3 values
+    """
+    i = np.int64(i)
+    j = np.int64(j)
+    k = np.int64(k)
+    m = np.int64(m)
+    o0 = (((a[0][0] * i) % m) + ((a[0][1] * j) % m) + ((a[0][2] * k) % m)) % m
+    o1 = (((a[1][0] * i) % m) + ((a[1][1] * j) % m) + ((a[1][2] * k) % m)) % m
+    o2 = (((a[2][0] * i) % m) + ((a[2][1] * j) % m) + ((a[2][2] * k) % m)) % m
+    return np.int32(o0), np.int32(o1), np.int32(o2)
+
+
+def _mrg_new_stream(seed=12345):
+    """
+    Create and return a new random state (vector of 6 32-bits integers).
+
+    Seed can be an integer or a vector of 6 integers.
+    :type seed: int | tuple
+    :return: a tuple with 6 32-bits integers.
+    """
+    if isinstance(seed, int):
+        if seed == 0:
+            raise Exception('seed should not be 0')
+        elif seed >= M2:
+            raise Exception('seed should be less than M2')
+        rstate = (np.int32(seed), np.int32(seed), np.int32(seed),
+                  np.int32(seed), np.int32(seed), np.int32(seed))
+    elif len(seed) == 6:
+        if seed[0] == 0 and seed[1] == 0 and seed[2] == 0:
+            raise Exception('First 3 values of seed should not be all 0')
+        if seed[3] == 0 and seed[4] == 0 and seed[5] == 0:
+            raise Exception('Last 3 values of seed should not be all 0')
+        if seed[0] >= M1 or seed[1] >= M1 or seed[2] >= M1:
+            raise Exception('First 3 values of seed should be less than M1')
+        if seed[3] >= M2 or seed[4] >= M2 or seed[5] >= M2:
+            raise Exception('Last 3 values of seed should be less than M2')
+        rstate = (np.int32(seed[0]), np.int32(seed[1]), np.int32(seed[2]),
+                  np.int32(seed[3]), np.int32(seed[4]), np.int32(seed[5]))
+    else:
+        raise Exception("seed should be 1 integer or 6 integers")
+    return rstate
+
+
+def _mrg_next_value(rstate, norm, mask, offset, dtype):
+    """
+    Generate next value from given MRG stream, norm, mask, offset and dtype.
+
+    Return new value and updated MRG stream.
+    """
+    x11 = rstate[0]
+    x12 = rstate[1]
+    x13 = rstate[2]
+    x21 = rstate[3]
+    x22 = rstate[4]
+    x23 = rstate[5]
+
+    y1 = (((x12 & MASK12) << i22) + (x12 >> i9) +
+          ((x13 & MASK13) << i7) + (x13 >> i24))
+
+    if y1 < 0 or y1 >= M1:
+        y1 = y1 - M1
+    y1 = y1 + x13
+    if y1 < 0 or y1 >= M1:
+        y1 = y1 - M1
+
+    x13 = x12
+    x12 = x11
+    x11 = y1
+
+    y1 = ((x21 & MASK2) << i15) + (MULT2 * (x21 >> i16))
+    if y1 < 0 or y1 >= M2:
+        y1 = y1 - M2
+    y2 = ((x23 & MASK2) << i15) + (MULT2 * (x23 >> i16))
+
+    if y2 < 0 or y2 >= M2:
+        y2 = y2 - M2
+    y2 = y2 + x23
+    if y2 < 0 or y2 >= M2:
+        y2 = y2 - M2
+    y2 = y2 + y1
+    if y2 < 0 or y2 >= M2:
+        y2 = y2 - M2
+
+    x23 = x22
+    x22 = x21
+    x21 = y2
+
+    new_rstate = x11, x12, x13, x21, x22, x23
+
+    if x11 <= x21:
+        value = SUPP_TYPES[dtype](((x11 - x21 + M1) & mask) + offset) * norm
+    else:
+        value = SUPP_TYPES[dtype](((x11 - x21) & mask) + offset) * norm
+    return value, new_rstate
+
+
+# ---------------
+# RNG public API.
+# ---------------
+
+
+@dataclass
+class MyiaRandomState:
+    """
+    Data class to represent a random state specialized for a d-type.
+
+    Supported dtypes: float16, float32, float64
+    """
+
+    rstate: tuple
+    mask: object
+    offset: object
+    norm: object
+    dtype: object
+
+
+def myia_random_state(dtype, seed=12345):
+    """
+    Generate a new random state for given dtype and seed.
+
+    :param dtype: a string: either 'float16', ´float32' or 'float64'.
+    :param seed: an integer or a tuple of 6 integers.
+    :return: a MyiaRandomState object.
+    """
+    rstate = _mrg_new_stream(seed)
+    if dtype == 'float16':
+        mask = 0x7fff
+        offset = 1
+        norm = np.float16(3.0458e-05)
+        dtype = FLOAT16
+    elif dtype == 'float32':
+        mask = 0xffffffff
+        offset = 0
+        norm = np.float32(4.6566126e-10)
+        dtype = FLOAT32
+    elif dtype == 'float64':
+        mask = 0xffffffff
+        offset = 0
+        norm = np.float64(4.656612873077392578125e-10)  # 1./2^31
+        dtype = FLOAT64
+    else:
+        raise Exception('Unsupported type for random state')
+    return MyiaRandomState(rstate, mask, offset, norm, dtype)
+
+
+def myia_increment_state(state: MyiaRandomState):
+    """
+    Increment a random state and return new updated random state.
+
+    Useful to create a new random state from a previous existing one.
+    Not useful in the random generation process.
+    :param state: an existing MyiaRandomState object.
+    :return: a MyiaRandomState object.
+    """
+    rs = state.rstate
+    new_rstate = (_mat_vec_mod_m(a_1_p_134, rs[0], rs[1], rs[2], M1) +
+                  _mat_vec_mod_m(a_2_p_134, rs[3], rs[4], rs[5], M2))
+    return MyiaRandomState(
+        new_rstate, state.mask, state.offset, state.norm, state.dtype)
+
+
+def myia_next_value(state: MyiaRandomState):
+    """
+    Generate a scalar using given state.
+
+    :param state: a MyiaRandomState object.
+    :return: a tuple: (new scalar value, updated state).
+    """
+    sample, new_rstate = _mrg_next_value(
+        state.rstate, state.norm, state.mask, state.offset, state.dtype)
+    return sample, MyiaRandomState(
+        new_rstate, state.mask, state.offset, state.norm, state.dtype)
+
+
+# --------------------------
+# MRG public high-level API.
+# --------------------------
+
+
+def myia_uniform(state: MyiaRandomState, low=0.0, high=1.0):
+    """
+    Generate a uniform scalar between low and high.
+
+    :param state: a MyiaRandomState object.
+    :param low: lower bound for uniform value.
+    :param high: upper bound for uniform value.
+    :return: a tuple: (new uniform scalar value, updated state).
+    """
+    u, next_state = myia_next_value(state)
+    r = u * (high - low) + low
+    return r, next_state

tests/test_rng.py

@@ -0,0 +1,89 @@
+""" Test random number generator.
+Notes:
+    - Compilation is slow (about 1 minute).
+    - Got several other compilation issues related to dtype handling, as dtype
+      is currently a string. We may need to set a better support for strings
+      in Myia.
+    - Exceptions are not correctly handled with @myia decorator.
+      Many "illegal" errors are raised, including illegal primitives,
+      illegal registered types, and other things. I can fix that,
+      but ultimately we will still need to correctly support strings
+      (so that exceptions could be correctly printed in backend).
+"""
+
+import numpy as np
+
+from myia import myia
+from myia.rng_mrg import (
+    MyiaRandomState,
+    myia_increment_state,
+    myia_random_state,
+    myia_uniform,
+)
+
+
+def pyth_random():
+    state = myia_random_state('float64')
+    v1, state = myia_uniform(state)
+    v2, state = myia_uniform(state)
+    v3, state = myia_uniform(state)
+    v4, state = myia_uniform(state)
+    v5, state = myia_uniform(state)
+    v6, state = myia_uniform(state)
+    v7, state = myia_uniform(state)
+    v8, state = myia_uniform(state)
+    v9, state = myia_uniform(state)
+    v10, _ = myia_uniform(state)
+    return v1, v2, v3, v4, v5, v6, v7, v8, v9, v10
+
+
+@myia
+def myia_random():
+    return pyth_random()
+
+
+@myia
+def run_myia_random_state():
+    return myia_random_state('float32')
+
+
+@myia
+def run_myia_increment_state(s):
+    return myia_increment_state(s)
+
+
+def test_myia_random_state():
+    state = run_myia_random_state()
+    assert isinstance(state, MyiaRandomState)
+    assert isinstance(state.rstate, tuple)
+    assert state.rstate == (np.int32(12345),) * 6
+
+
+def test_myia_increment_state():
+    state_1 = run_myia_random_state()
+    state_2 = run_myia_increment_state(state_1)
+    state_3 = run_myia_increment_state(state_2)
+    state_4 = run_myia_increment_state(state_3)
+    assert state_2.rstate == tuple(np.int32(val) for val in (
+        336690377, 597094797, 1245771585, 85196284, 523477687, 2094976052))
+    assert state_3.rstate == tuple(np.int32(val) for val in (
+        502033783, 1322587635, 1964121530, 1949818481, 1607232546, 1462898381))
+    assert state_4.rstate == tuple(np.int32(val) for val in (
+        739421137, 1475938232, 730262207, 1630192198, 324551134, 795289868))
+
+
+def test_myia_random_generation():
+    # Hide runtime warnings about overflow in integer operations.
+    err_orig = np.seterr(all='ignore')
+
+    pyth_res = pyth_random()
+    myia_res = myia_random()
+    assert pyth_res == myia_res
+    assert pyth_res == tuple(np.float64(val) for val in (
+        0.7353244530968368, 0.6142074400559068, 0.11007806099951267,
+        0.6487741703167558, 0.36619443260133266, 0.10882294131442904,
+        0.5330547927878797, 0.9783797566778958, 0.9151237849146128,
+        0.8509745532646775))
+
+    # Get numpy back to its original warning config.
+    np.seterr(**err_orig)