def _benchmarkUniform(self, name, dtype, use_xla_jit):
def builder_fn():
shape = (10, 1000, 1000)
seed_var = variables.Variable((312, 456),
dtype=dtypes.int32,
name='input')
random_t = stateless.stateless_random_uniform(
shape, seed=seed_var, dtype=dtype)
return '%s.shape%s' % (name, shape), [random_t]
xla_test.Benchmark(self, builder_fn, use_xla_jit=use_xla_jit, device='cpu')
def _benchmarkGather(self, name, axis, gather_indices, use_xla_jit):
def BuilderFn():
inputs = variables.Variable(
array_ops.zeros([100, 100, 10, 100, 50], dtype=dtypes.float32),
dtype=dtypes.float32,
name='input')
indices = variables.Variable(
gather_indices, dtype=dtypes.int32, name='indices')
gather_t = array_ops.gather(inputs, indices, axis=axis)
return '%s.axis%d' % (name, axis), [gather_t]
xla_test.Benchmark(self, BuilderFn, use_xla_jit=use_xla_jit, device='cpu')
def benchmarkLayerTrainingXLA(self):
xla_test.Benchmark(self, lambda: self._LayerBuilder(True), True,
FLAGS.device)
def benchmarkLayerInferenceXLA(self):
xla_test.Benchmark(self, lambda: self._LayerBuilder(False), True,
FLAGS.device)
