def _sample_n(self, n, seed=None):
seed = samplers.sanitize_seed(seed)
seed1, seed2 = samplers.split_seed(seed, salt='Skellam')
log_rate1 = self._log_rate1_parameter_no_checks()
log_rate2 = self._log_rate2_parameter_no_checks()
batch_shape = self._batch_shape_tensor(
log_rate1=log_rate1, log_rate2=log_rate2)
log_rate1 = ps.broadcast_to(log_rate1, batch_shape)
log_rate2 = ps.broadcast_to(log_rate2, batch_shape)
sample1 = poisson_lib.random_poisson(
[n], log_rates=log_rate1, seed=seed1)[0]
sample2 = poisson_lib.random_poisson(
[n], log_rates=log_rate2, seed=seed2)[0]
return sample1 - sample2
def testSampleCPU(self):
with tf.device('CPU'):
_, runtime = self.evaluate(
poisson_lib.random_poisson(shape=tf.constant([],
dtype=tf.int32),
rates=tf.constant(10.),
seed=test_util.test_seed()))
self.assertEqual(implementation_selection._RUNTIME_CPU, runtime)
def testSampleGPU(self):
if not tf.test.is_gpu_available():
self.skipTest('no GPU')
with tf.device('GPU'):
_, runtime = self.evaluate(poisson_lib.random_poisson(
shape=tf.constant([], dtype=tf.int32),
rates=tf.constant(10.),
seed=test_util.test_seed()))
self.assertEqual(implementation_selection._RUNTIME_DEFAULT, runtime)