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)