message: Add benchmark for QR decomposition in TensorFlow.
diff: @@ -21,13 +21,18 @@ from __future__ import print_function
  import numpy as np
  from tensorflow.python import tf2
  +from tensorflow.python.client import session
  from tensorflow.python.framework import constant_op
  +from tensorflow.python.framework import ops
  from tensorflow.python.framework import test_util
  from tensorflow.python.ops import array_ops
  +from tensorflow.python.ops import control_flow_ops
  from tensorflow.python.ops import gradient_checker
  from tensorflow.python.ops import linalg_ops
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops import random_ops
  +from tensorflow.python.ops import variables
  +from tensorflow.python.platform import benchmark
  from tensorflow.python.platform import test
  @@ -200,6 +205,60 @@ def _GetQrGradOpTest(dtype_, shape_, full_matrices_):
  return Test
  +class QRBenchmark(test.Benchmark):
  +
  + shapes = [
  + (4, 4),
  + (8, 8),
  + (16, 16),
  + (101, 101),
  + (256, 256),
  + (1024, 1024),
  + (2048, 2048),
  + (1024, 2),
  + (1024, 32),
  + (1024, 128),
  + (1024, 512),
  + (1, 8, 8),
  + (10, 8, 8),
  + (100, 8, 8),
  + (1, 256, 256),
  + (10, 256, 256),
  + (100, 256, 256),
  + ]
  +
  + def benchmarkQROp(self):
  + for shape_ in self.shapes:
  + with ops.Graph().as_default(), \
  + session.Session(config=benchmark.benchmark_config()) as sess, \
  + ops.device("/cpu:0"):
  + matrix_value = np.random.uniform(
  + low=-1.0, high=1.0, size=shape_).astype(np.float32)
  + matrix = variables.Variable(matrix_value)
  + q, r = linalg_ops.qr(matrix)
  + variables.global_variables_initializer().run()
  + self.run_op_benchmark(
  + sess,
  + control_flow_ops.group(q, r),
  + min_iters=25,
  + name="QR_cpu_{shape}".format(shape=shape_))
  +
  + if test.is_gpu_available(True):
  + with ops.Graph().as_default(), \
  + session.Session(config=benchmark.benchmark_config()) as sess, \
  + ops.device("/device:GPU:0"):
  + matrix_value = np.random.uniform(
  + low=-1.0, high=1.0, size=shape_).astype(np.float32)
  + matrix = variables.Variable(matrix_value)
  + q, r = linalg_ops.qr(matrix)
  + variables.global_variables_initializer().run()
  + self.run_op_benchmark(
  + sess,
  + control_flow_ops.group(q, r),
  + min_iters=25,
  + name="QR_gpu_{shape}".format(shape=shape_))
  +
  +
  if __name__ == "__main__":
  for dtype in np.float32, np.float64, np.complex64, np.complex128:
  for rows in 1, 2, 5, 10, 32, 100:
  