def pytest_generate_tests(metafunc):
if 'test_bijection' in metafunc.funcargnames:
vals = []
ids = []
for name, words, bitness in itertools.product(['threefry', 'philox'],
[2, 4], [32, 64]):
val = TestBijection(name, words, bitness)
vals.append(val)
ids.append(str(val))
metafunc.parametrize('test_bijection', vals, ids=ids)
if 'test_sampler_int' in metafunc.funcargnames:
bijection = philox(64, 4)
vals = [TestUniformInteger(bijection, -10, 98)]
ids = [test.name for test in vals]
metafunc.parametrize('test_sampler_int', vals, ids=ids)
if 'test_sampler_float' in metafunc.funcargnames:
bijection = philox(64, 4)
vals = [
TestUniformFloat(bijection, -5, 7.7),
TestNormalBM(bijection, -2, 10),
TestNormalBMComplex(bijection, -3 + 4j, 7),
TestGamma(bijection, 3, 10)
]
ids = [test.name for test in vals]
metafunc.parametrize('test_sampler_float', vals, ids=ids)
def pytest_generate_tests(metafunc):
if 'test_bijection' in metafunc.funcargnames:
vals = []
ids = []
for name, words, bitness in itertools.product(['threefry', 'philox'], [2, 4], [32, 64]):
val = TestBijection(name, words, bitness)
vals.append(val)
ids.append(str(val))
metafunc.parametrize('test_bijection', vals, ids=ids)
if 'test_sampler_int' in metafunc.funcargnames:
bijection = philox(64, 4)
vals = [
TestUniformInteger(bijection, -10, 98)]
ids = [test.name for test in vals]
metafunc.parametrize('test_sampler_int', vals, ids=ids)
if 'test_sampler_float' in metafunc.funcargnames:
bijection = philox(64, 4)
vals = [
TestUniformFloat(bijection, -5, 7.7),
TestNormalBM(bijection, -2, 10),
TestNormalBMComplex(bijection, -3 + 4j, 7),
TestGamma(bijection, 3, 10)]
ids = [test.name for test in vals]
metafunc.parametrize('test_sampler_float', vals, ids=ids)
def test_64_to_32_bit(thr):
extent = (0, 2**31 - 1)
mean, std = uniform_discrete_mean_and_std(*extent)
bijection = philox(64, 4)
sampler = uniform_integer(bijection, numpy.uint32, extent[0],
extent[1] + 1)
check_kernel_sampler(thr, sampler, extent=extent, mean=mean, std=std)
def _build_plan(self, plan_factory, device_params, alpha, beta, seed):
plan = plan_factory()
bijection = philox(64, 2)
# Keeping the kernel the same so it can be cached.
# The seed will be passed as the computation parameter instead.
keygen = KeyGenerator.create(bijection, seed=numpy.int32(0))
sampler = normal_bm(bijection, numpy.float64)
squeezing = plan.persistent_array(self._system.squeezing)
decoherence = plan.persistent_array(self._system.decoherence)
plan.kernel_call(TEMPLATE.get_def("generate_input_state"),
[alpha, beta, squeezing, decoherence, seed],
kernel_name="generate",
global_size=alpha.shape,
render_kwds=dict(
system=self._system,
representation=self._representation,
Representation=Representation,
bijection=bijection,
keygen=keygen,
sampler=sampler,
ordering=ordering,
exp=functions.exp(numpy.float64),
mul_cr=functions.mul(numpy.complex128,
numpy.float64),
add_cc=functions.add(numpy.complex128,
numpy.complex128),
))
return plan
def test_computation_performance(thr_and_double, fast_math, test_sampler_float):
thr, double = thr_and_double
size = 2 ** 15
batch = 2 ** 6
bijection = philox(64, 4)
sampler = test_sampler_float.get_sampler(bijection, double)
rng = CBRNG(Type(sampler.dtype, shape=(batch, size)), 1, sampler)
dest_dev = thr.empty_like(rng.parameter.randoms)
counters = rng.create_counters()
counters_dev = thr.to_device(counters)
rngc = rng.compile(thr, fast_math=fast_math)
attempts = 10
times = []
for i in range(attempts):
t1 = time.time()
rngc(counters_dev, dest_dev)
thr.synchronize()
times.append(time.time() - t1)
byte_size = size * batch * sampler.dtype.itemsize
return min(times), byte_size
def __call__(self,
cls,
randoms_arr,
generators_dim,
sampler_kwds=None,
seed=None):
bijection = philox(64, 4)
if sampler_kwds is None:
sampler_kwds = {}
sampler = self._sampler_func(bijection, randoms_arr.dtype,
**sampler_kwds)
return cls(randoms_arr, generators_dim, sampler, seed=seed)
def test_computation_general(thr_and_double):
size = 10000
batch = 101
thr, double = thr_and_double
dtype = numpy.float64 if double else numpy.float32
mean, std = -2, 10
bijection = philox(64, 4)
sampler = normal_bm(bijection, dtype, mean=mean, std=std)
rng = CBRNG(Type(dtype, shape=(batch, size)), 1, sampler)
check_computation(thr, rng, mean=mean, std=std)
def test_computation_general(thr_and_double):
size = 10000
batch = 101
thr, double = thr_and_double
bijection = philox(64, 4)
ref = NormalBMHelper(mean=-2, std=10)
sampler = ref.get_sampler(bijection, double)
rng = CBRNG(Type(sampler.dtype, shape=(batch, size)), 1, sampler)
check_computation(thr, rng, ref)
def test_computation_general(thr_and_double):
size = 10000
batch = 101
thr, double = thr_and_double
dtype = numpy.float64 if double else numpy.float32
mean, std = -2, 10
bijection = philox(64, 4)
sampler = normal_bm(bijection, dtype, mean=mean, std=std)
rng = CBRNG(Type(dtype, shape=(batch, size)), 1, sampler)
check_computation(thr, rng, mean=mean, std=std)
def _build_plan(self, plan_factory, device_params, alpha, beta, alpha_i,
beta_i, seed):
plan = plan_factory()
system = self._system
representation = self._representation
unitary = plan.persistent_array(self._system.unitary)
needs_noise_matrix = representation != Representation.POSITIVE_P and system.needs_noise_matrix(
)
mmul = MatrixMul(alpha, unitary, transposed_b=True)
if not needs_noise_matrix:
# TODO: this could be sped up for repr != POSITIVE_P,
# since in that case alpha == conj(beta), and we don't need to do two multuplications.
mmul_beta = MatrixMul(beta, unitary, transposed_b=True)
trf_conj = self._make_trf_conj()
mmul_beta.parameter.matrix_b.connect(trf_conj,
trf_conj.output,
matrix_b_p=trf_conj.input)
plan.computation_call(mmul, alpha, alpha_i, unitary)
plan.computation_call(mmul_beta, beta, beta_i, unitary)
else:
noise_matrix = system.noise_matrix()
noise_matrix_dev = plan.persistent_array(noise_matrix)
# If we're here, it's not positive-P, and alpha == conj(beta).
# This means we can just calculate alpha, and then build beta from it.
w = plan.temp_array_like(alpha)
temp_alpha = plan.temp_array_like(alpha)
plan.computation_call(mmul, temp_alpha, alpha_i, unitary)
bijection = philox(64, 2)
# Keeping the kernel the same so it can be cached.
# The seed will be passed as the computation parameter instead.
keygen = KeyGenerator.create(bijection, seed=numpy.int32(0))
sampler = normal_bm(bijection, numpy.float64)
plan.kernel_call(TEMPLATE.get_def("generate_apply_matrix_noise"),
[w, seed],
kernel_name="generate_apply_matrix_noise",
global_size=alpha.shape,
render_kwds=dict(
bijection=bijection,
keygen=keygen,
sampler=sampler,
mul_cr=functions.mul(numpy.complex128,
numpy.float64),
add_cc=functions.add(numpy.complex128,
numpy.complex128),
))
noise = plan.temp_array_like(alpha)
plan.computation_call(mmul, noise, w, noise_matrix_dev)
plan.kernel_call(TEMPLATE.get_def("add_noise"),
[alpha, beta, temp_alpha, noise],
kernel_name="add_noise",
global_size=alpha.shape,
render_kwds=dict(
add=functions.add(numpy.complex128,
numpy.complex128),
conj=functions.conj(numpy.complex128)))
return plan
def test_kernel_sampler_float(thr_and_double, test_sampler_float):
thr, double = thr_and_double
bijection = philox(64, 4)
check_kernel_sampler(
thr, test_sampler_float.get_sampler(bijection, double), test_sampler_float)
def test_kernel_sampler_int(thr, test_sampler_int):
bijection = philox(64, 4)
check_kernel_sampler(
thr, test_sampler_int.get_sampler(bijection, numpy.int32), test_sampler_int)
def test_64_to_32_bit(thr):
bijection = philox(64, 4)
ref = UniformIntegerHelper(0, 2**31-1)
sampler = ref.get_sampler(bijection, numpy.uint32)
check_kernel_sampler(thr, sampler, ref)
def test_32_to_64_bit(thr):
bijection = philox(32, 4)
ref = UniformIntegerHelper(0, 2**63-1)
sampler = ref.get_sampler(bijection, numpy.uint64)
check_kernel_sampler(thr, sampler, ref)
def __call__(self, cls, randoms_arr, generators_dim, sampler_kwds=None, seed=None):
bijection = philox(64, 4)
if sampler_kwds is None:
sampler_kwds = {}
sampler = self._sampler_func(bijection, randoms_arr.dtype, **sampler_kwds)
return cls(randoms_arr, generators_dim, sampler, seed=seed)
def test_64_to_32_bit(thr):
extent = (0, 2**31-1)
mean, std = uniform_discrete_mean_and_std(*extent)
bijection = philox(64, 4)
sampler = uniform_integer(bijection, numpy.uint32, extent[0], extent[1] + 1)
check_kernel_sampler(thr, sampler, extent=extent, mean=mean, std=std)