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)