def prng(s):
"""Creates RNG state from seed.
This implementation doesn't pass RNG states explicitly so the result is
always a dummy 0.
Args:
s: the seed, an integer.
Returns:
A dummy integer 0.
"""
# TODO(wangpeng): change it to use stateless random ops to truely mimic JAX
# RNGs
random.seed(s)
# Returning None will cause errors in some layer/optimizer libraries based on
# JAX
return asarray(0, dtype=np.int64)
def run_test(*args):
num_samples = 1000
tol = 0.1 # High tolerance to keep the # of samples low else the test
# takes a long time to run.
random.seed(10)
outputs = [random.randn(*args) for _ in range(num_samples)]
# Test output shape.
for output in outputs:
self.assertEqual(output.shape, tuple(args))
self.assertEqual(output.dtype.type, random.DEFAULT_RANDN_DTYPE)
if np.prod(args): # Don't bother with empty arrays.
outputs = [output.tolist() for output in outputs]
# Test that the properties of normal distribution are satisfied.
mean = np.mean(outputs, axis=0)
stddev = np.std(outputs, axis=0)
self.assertAllClose(mean, np.zeros(args), atol=tol)
self.assertAllClose(stddev, np.ones(args), atol=tol)
# Test that outputs are different with different seeds.
random.seed(20)
diff_seed_outputs = [
random.randn(*args).tolist() for _ in range(num_samples)
]
self.assertNotAllClose(outputs, diff_seed_outputs)
# Test that outputs are the same with the same seed.
random.seed(10)
same_seed_outputs = [
random.randn(*args).tolist() for _ in range(num_samples)
]
self.assertAllClose(outputs, same_seed_outputs)