Fix pylint errors in benchmark directory

benchmarks/scripts/benchmark_clifford_circuit.py

@@ -23,15 +23,17 @@
 
 from tensorflow_quantum.core.ops import tfq_simulate_ops
 from tensorflow_quantum.core.serialize.serializer import serialize_circuit
-from models.random_clifford_circuit import random_clifford_circuit
-import flags
-import benchmark_util
+from benchmarks.scripts.models.random_clifford_circuit import (
+    random_clifford_circuit)
+
+from benchmarks.scripts import flags
+from benchmarks.scripts import benchmark_util
 
 SEED = 48510234
 SRC = os.path.dirname(os.path.realpath(__file__))
 os.environ['TEST_REPORT_FILE_PREFIX'] = os.path.join(SRC, 'reports/')
-TEST_PARAMS_1 = flags.TEST_FLAGS(n_qubits=3, n_moments=5, op_density=0.99)
-TEST_PARAMS_2 = flags.TEST_FLAGS(n_qubits=4, n_moments=5, op_density=0.99)
+TEST_PARAMS_1 = flags.test_flags(n_qubits=3, n_moments=5, op_density=0.99)
+TEST_PARAMS_2 = flags.test_flags(n_qubits=4, n_moments=5, op_density=0.99)
 ALL_PARAMS = [TEST_PARAMS_1, TEST_PARAMS_2]
 
 
@@ -42,7 +44,7 @@ class CliffordBenchmarksTest(tf.test.TestCase, parameterized.TestCase):
         ("params_1", TEST_PARAMS_1),
         ("params_2", TEST_PARAMS_2),
     )
-    def testBenchmarkCliffordCircuitEager(self, params):
+    def test_benchmark_clifford_circuit_eager(self, params):
         """Test that Op constructs and runs correctly."""
         proto_file_path = os.path.join(
             SRC, "reports/",
@@ -77,7 +79,7 @@ class CliffordBenchmarks(tf.test.Benchmark):
 
     def __init__(self, params=None):
         """Pull in command line flags or use provided flags."""
-        super(CliffordBenchmarks, self).__init__()
+        super().__init__()
         # Allow input params for testing purposes.
         self.params = params if params else flags.FLAGS
 

benchmarks/scripts/benchmark_op_gradients.py

@@ -23,22 +23,23 @@
 import numpy as np
 
 from tensorflow_quantum.core.ops import tfq_simulate_ops
-import benchmark_util
-import flags
 
 from tensorflow_quantum.python import util
 from tensorflow_quantum.python.differentiators import (
     linear_combination,
     parameter_shift,
 )
 
+from benchmarks.scripts import benchmark_util
+from benchmarks.scripts import flags
+
 SRC = os.path.dirname(os.path.realpath(__file__))
 os.environ['TEST_REPORT_FILE_PREFIX'] = os.path.join(SRC, 'reports/')
-TEST_PARAMS_1 = flags.TEST_FLAGS(n_symbols=4,
+TEST_PARAMS_1 = flags.test_flags(n_symbols=4,
                                  n_qubits=3,
                                  n_moments=5,
                                  op_density=0.9)
-TEST_PARAMS_2 = flags.TEST_FLAGS(n_symbols=3,
+TEST_PARAMS_2 = flags.test_flags(n_symbols=3,
                                  n_qubits=4,
                                  n_moments=5,
                                  op_density=0.6)
@@ -58,7 +59,7 @@ class GradientBenchmarksTest(tf.test.TestCase, parameterized.TestCase):
                     ],
                     'params': [TEST_PARAMS_1, TEST_PARAMS_2]
                 })))
-    def testBenchmarkGradient(self, diff, params):
+    def test_benchmark_gradient(self, diff, params):
         """Test that op constructs and runs correctly."""
 
         bench_name = "GradientBenchmarks.{}_{}_{}_{}_{}".format(
@@ -93,7 +94,7 @@ class GradientBenchmarks(tf.test.Benchmark):
 
     def __init__(self, params=None):
         """Pull in command line flags or use provided flags."""
-        super(GradientBenchmarks, self).__init__()
+        super().__init__()
         self.params = params if params else flags.FLAGS
         self.setup()
 
@@ -110,7 +111,8 @@ def setup(self):
                                  replace=True)))
         symbol_names = list(symbol_names)
 
-        circuit_batch, resolver_batch = util.random_symbol_circuit_resolver_batch(
+        circuit_batch, resolver_batch = \
+            util.random_symbol_circuit_resolver_batch(
             qubits=qubits,
             symbols=symbol_names,
             batch_size=self.params.batch_size,

benchmarks/scripts/benchmark_random_circuit.py

@@ -23,14 +23,14 @@
 
 from tensorflow_quantum.core.ops import tfq_simulate_ops
 from tensorflow_quantum.core.serialize.serializer import serialize_circuit
-import flags
-import benchmark_util
+from benchmarks.scripts import flags
+from benchmarks.scripts import benchmark_util
 
 SEED = 63536323
 SRC = os.path.dirname(os.path.realpath(__file__))
 os.environ['TEST_REPORT_FILE_PREFIX'] = os.path.join(SRC, 'reports/')
-TEST_PARAMS_1 = flags.TEST_FLAGS(n_rows=3, n_cols=5, n_moments=5)
-TEST_PARAMS_2 = flags.TEST_FLAGS(n_rows=4, n_cols=4, n_moments=20)
+TEST_PARAMS_1 = flags.test_flags(n_rows=3, n_cols=5, n_moments=5)
+TEST_PARAMS_2 = flags.test_flags(n_rows=4, n_cols=4, n_moments=20)
 
 
 def make_random_circuit(n_rows, n_cols, depth):
@@ -49,7 +49,7 @@ class RandomCircuitBenchmarksTest(tf.test.TestCase, parameterized.TestCase):
         ("params_1", TEST_PARAMS_1),
         ("params_2", TEST_PARAMS_2),
     )
-    def testBenchmarkRandomCircuit(self, params):
+    def test_benchmark_random_circuit(self, params):
         """Test that Op constructs and runs correctly."""
         proto_file_path = os.path.join(
             SRC, "reports/",
@@ -77,7 +77,7 @@ def testBenchmarkRandomCircuit(self, params):
         ("params_1", TEST_PARAMS_1),
         ("params_2", TEST_PARAMS_2),
     )
-    def testRandomCircuitParams(self, params):
+    def test_random_circuit_params(self, params):
         """Ensure that the random circuits are structured as advertised."""
         circuit = make_random_circuit(params.n_rows, params.n_cols,
                                       params.n_moments)
@@ -95,7 +95,7 @@ class RandomCircuitBenchmarks(tf.test.Benchmark):
 
     def __init__(self, params=None):
         """Pull in command line flags or use provided flags."""
-        super(RandomCircuitBenchmarks, self).__init__()
+        super().__init__()
         # Allow input params for testing purposes.
         self.params = params if params else flags.FLAGS
 
@@ -106,7 +106,8 @@ def _simulate_circuit(self, circuit, params):
             [[0]] * params.batch_size)
 
     def benchmark_random_circuit(self):
-        """Benchmark simulator performance on a classically intractable circuit."""
+        """Benchmark simulator performance on
+        a classically intractable circuit."""
 
         circuit = make_random_circuit(self.params.n_rows, self.params.n_cols,
                                       self.params.n_moments)

benchmarks/scripts/benchmark_util.py

@@ -14,10 +14,18 @@
 # ==============================================================================
 """Utility functions for benchmark tools."""
 import tensorflow as tf
-import test_log_pb2
+from tensorflow.core.util import test_log_pb2
 
 
 def read_benchmark_entry(f):
+    """Reads a benchmark entry from a file.
+
+    Args:
+        f: File path to read from.
+
+    Returns:
+        The first entry in the benchmark file.
+    """
     s = tf.io.gfile.GFile(f, "rb").read()
     entries = test_log_pb2.BenchmarkEntries.FromString(s)
     return entries.entry[0]

benchmarks/scripts/benchmark_util_test.py

@@ -18,8 +18,8 @@
 
 import tensorflow as tf
 
-import test_log_pb2
-import benchmark_util
+from tensorflow.core.util import test_log_pb2
+from benchmarks.scripts import benchmark_util
 
 
 def _make_dummy_benchmark_report():

benchmarks/scripts/flags.py

@@ -84,7 +84,7 @@
                           lower_bound=0)
 
 
-def TEST_FLAGS(**kwargs):
+def test_flags(**kwargs):
     """Create a set of test flags by kwarg assignment.
 
     This constructs a named tuple that mimics the interface of absl.flags.

benchmarks/scripts/models/random_clifford_circuit.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
+"""Module for generating random Clifford circuits."""
 
 from typing import Iterable
 

benchmarks/scripts/models/random_clifford_circuit_test.py

@@ -12,13 +12,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
+"""Tests for the random Clifford circuit generator."""
 
 from absl.testing import parameterized
 import cirq
 import numpy as np
 import tensorflow as tf
 
-from random_clifford_circuit import random_clifford_circuit
+from benchmarks.scripts.models.random_clifford_circuit import (
+    random_clifford_circuit)
 
 
 class RandomCliffordCircuitTest(parameterized.TestCase, tf.test.TestCase):