def testSampleCPU(self):
with tf.device('CPU'):
_, runtime = self.evaluate(
binomial_lib._random_binomial(
shape=tf.constant([], dtype=tf.int32),
counts=tf.constant(10.), probs=tf.constant(.5),
seed=test_util.test_seed()))
self.assertEqual(implementation_selection._RUNTIME_CPU, runtime)
def testSampleGPU(self):
# Debugging tip: --vmodule=implementation_selector=2,function_api_info=3
# helps debug errors like: Skipping optimization due to error while loading
# function libraries: Invalid argument: Functions
# '__inference__binomial_cpu_1184' and '__inference__binomial_noncpu_2305'
# both implement 'binomial_76a77701-978a-4299-ad23-ff0d5c7598cb' but their
# signatures do not match.
if not tf.test.is_gpu_available():
self.skipTest('no GPU')
with tf.device('GPU'):
_, runtime = self.evaluate(binomial_lib._random_binomial(
shape=tf.constant([], dtype=tf.int32),
counts=tf.constant(10.), probs=tf.constant(.5),
seed=test_util.test_seed()))
self.assertEqual(implementation_selection._RUNTIME_DEFAULT, runtime)
