Add benchmarks for matrix_inverse and determinant.

Add missing op grouping in Cholesky benchmark.

PiperOrigin-RevId: 164281947

1 files changed, 50 insertions, 0 deletions

diff --git a/tensorflow/python/kernel_tests/determinant_op_test.py b/tensorflow/python/kernel_tests/determinant_op_test.py
index a4daefec39..b9fc110405 100644
--- a/tensorflow/python/kernel_tests/determinant_op_test.py
+++ b/tensorflow/python/kernel_tests/determinant_op_test.py
@@ -20,7 +20,10 @@ from __future__ import print_function
 
 import numpy as np
 
+from tensorflow.python.client import session
 from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import control_flow_ops
 from tensorflow.python.ops import linalg_ops
 from tensorflow.python.platform import test
 
@@ -125,5 +128,52 @@ class DeterminantOpTest(test.TestCase):
     self._compareDeterminant(np.empty([2, 0, 0]))
 
 
+class MatrixDeterminantBenchmark(test.Benchmark):
+
+  sizes = [
+      (4, 4),
+      (16, 16),
+      (256, 256),
+      (1024, 1024),
+      (513, 4, 4),
+      (513, 16, 16),
+      (513, 256, 256),
+  ]
+
+  def _GenerateData(self, size):
+    batch_shape = size[:-2]
+    size = size[-2:]
+    assert size[0] == size[1]
+    n = size[0]
+    data = np.ones(size).astype(np.float32) / (
+        2.0 * n) + np.diag(np.ones(n).astype(np.float32))
+    return np.tile(data, batch_shape + (1, 1))
+
+  def benchmarkMatrixDeterminantOp(self):
+    for size in self.sizes:
+      data = self._GenerateData(size)
+
+      with ops.Graph().as_default(), session.Session() as sess, ops.device(
+          "/cpu:0"):
+        d = linalg_ops.matrix_determinant(data)
+        self.run_op_benchmark(
+            sess,
+            control_flow_ops.group(
+                d,),
+            min_iters=25,
+            name="matrix_determinant_cpu_{size}".format(size=size))
+
+      if test.is_gpu_available(True):
+        with ops.Graph().as_default(), session.Session() as sess, ops.device(
+            "/gpu:0"):
+          d = linalg_ops.matrix_determinant(data)
+          self.run_op_benchmark(
+              sess,
+              control_flow_ops.group(
+                  d,),
+              min_iters=25,
+              name="matrix_determinant_gpu_{size}".format(size=size))
+
+
 if __name__ == "__main__":
   test.main()