uv tool run yapf -irp .

Author: cclauss

compiler_opt/benchmark/benchmark_chromium.py

@@ -206,7 +206,9 @@ def main(_):
     with open(test_description, encoding='UTF-8') as test_description_file:
       print(test_description)
       test_descriptions.append(json.load(test_description_file))
-  test_executables = [test_description['executable'] for test_description in test_descriptions]
+  test_executables = [
+      test_description['executable'] for test_description in test_descriptions
+  ]
 
   if FLAGS.compile_llvm:
     benchmarking_utils.build_llvm(FLAGS.model_path, FLAGS.llvm_use_incremental,

compiler_opt/benchmark/benchmark_report_test.py

@@ -82,8 +82,7 @@ def test_loading(self):
             }
         })
     self.assertSetEqual(report.names(), {'BM_A', 'BM_B'})
-    self.assertSetEqual(report.counters(),
-                        {'PerfCounter_0', 'PerfCounter_1'})
+    self.assertSetEqual(report.counters(), {'PerfCounter_0', 'PerfCounter_1'})
     self.assertEqual(
         report.counter_means('BM_A', 'PerfCounter_0'),
         (10.488088481701517, 0.7071067811865476))

compiler_opt/es/blackbox_learner.py

@@ -170,8 +170,9 @@ def _get_perturbations(self) -> List[npt.NDArray[np.float32]]:
     """Get perturbations for the model weights."""
     rng = np.random.default_rng(seed=self._seed)
     return [
-      rng.normal(size=len(self._model_weights)) * self._config.precision_parameter
-      for _ in range(self._config.total_num_perturbations)
+        rng.normal(size=len(self._model_weights)) *
+        self._config.precision_parameter
+        for _ in range(self._config.total_num_perturbations)
     ]
 
   def _update_model(self, perturbations: List[npt.NDArray[np.float32]],
@@ -274,8 +275,10 @@ def run_step(self, pool: FixedWorkerPool) -> None:
           p for p in initial_perturbations for p in (p, -p)
       ]
 
-    perturbations_as_policies = [self._get_policy_from_perturbation(perturbation)
-                                 for perturbation in initial_perturbations]
+    perturbations_as_policies = [
+        self._get_policy_from_perturbation(perturbation)
+        for perturbation in initial_perturbations
+    ]
 
     results = self._evaluator.get_results(pool, perturbations_as_policies)
     rewards = self._evaluator.get_rewards(results)

compiler_opt/es/regalloc_trace/regalloc_trace_worker.py

@@ -113,9 +113,9 @@ def _build_corpus(self, modules: Collection[corpus.ModuleSpec],
     with concurrent.futures.ThreadPoolExecutor(
         max_workers=self._thread_count) as thread_pool:
       compile_futures = [
-        thread_pool.submit(self._compile_module, module, output_directory,
-                           tflite_policy_dir) for module in modules
-        ]
+          thread_pool.submit(self._compile_module, module, output_directory,
+                             tflite_policy_dir) for module in modules
+      ]
 
       for future in compile_futures:
         if future.exception() is not None:

compiler_opt/rl/data_reader.py

@@ -45,9 +45,11 @@ def _parser_fn(serialized_proto):
     context_features = {}
     # pylint: disable=g-complex-comprehension
     sequence_features = {
-        tensor_spec.name: tf.io.FixedLenSequenceFeature(
-             shape=tensor_spec.shape, dtype=tensor_spec.dtype)
-        for tensor_spec in agent_cfg.time_step_spec.observation.values()}
+        tensor_spec.name:
+            tf.io.FixedLenSequenceFeature(
+                shape=tensor_spec.shape, dtype=tensor_spec.dtype)
+        for tensor_spec in agent_cfg.time_step_spec.observation.values()
+    }
     sequence_features[
         agent_cfg.action_spec.name] = tf.io.FixedLenSequenceFeature(
             shape=agent_cfg.action_spec.shape,

compiler_opt/rl/imitation_learning/weighted_bc_trainer_lib.py

@@ -235,9 +235,11 @@ def __init__(
 
     observation_spec, action_spec = config.get_inlining_signature_spec()
     sequence_features = {
-        tensor_spec.name: tf.io.FixedLenSequenceFeature(
-             shape=tensor_spec.shape, dtype=tensor_spec.dtype)
-        for tensor_spec in observation_spec[-1].values()}
+        tensor_spec.name:
+            tf.io.FixedLenSequenceFeature(
+                shape=tensor_spec.shape, dtype=tensor_spec.dtype)
+        for tensor_spec in observation_spec[-1].values()
+    }
     sequence_features.update({
         action_spec.name:
             tf.io.FixedLenSequenceFeature(

compiler_opt/rl/inlining/config.py

@@ -64,7 +64,8 @@ def get_inlining_signature_spec():
           'is_multiple_blocks',
           'nested_inlines',
           'nested_inline_cost_estimate',
-          'threshold')}
+          'threshold')
+  }
   reward_spec = tf.TensorSpec(dtype=tf.float32, shape=(), name='reward')
   time_step_spec = time_step.time_step_spec(observation_spec, reward_spec)
   action_spec = tensor_spec.BoundedTensorSpec(

compiler_opt/rl/inlining/imitation_learning_config.py

@@ -33,14 +33,15 @@ def get_inlining_signature_spec():
   """Returns (time_step_spec, action_spec) for collecting IL trajectories."""
   time_step_spec, _ = config.get_inlining_signature_spec()
   observation_spec = time_step_spec.observation
-  observation_spec.update(
-      {key: tf.TensorSpec(dtype=tf.int64, shape=(), name=key) for key in (
+  observation_spec.update({
+      key: tf.TensorSpec(dtype=tf.int64, shape=(), name=key) for key in (
           'is_callee_avail_external',
           'is_caller_avail_external',
           # following features are not used in training.
           'inlining_default',
           SequenceExampleFeatureNames.label_name,
-          'policy_label')})  # testing only
+          'policy_label')
+  })  # testing only
 
   observation_spec[SequenceExampleFeatureNames.module_name] = tf.TensorSpec(
       dtype=tf.string, shape=(), name=SequenceExampleFeatureNames.module_name)
@@ -63,11 +64,12 @@ def get_input_signature():
   """Returns (time_step_spec, action_spec) wrapping a trained policy."""
   time_step_spec, action_spec = config.get_inlining_signature_spec()
   observation_spec = time_step_spec.observation
-  observation_spec.update(
-      {key: tf.TensorSpec(dtype=tf.int64, shape=(), name=key) for key in (
+  observation_spec.update({
+      key: tf.TensorSpec(dtype=tf.int64, shape=(), name=key) for key in (
           'is_callee_avail_external',
           'is_caller_avail_external',
-      )})
+      )
+  })
 
   time_step_spec = time_step.time_step_spec(observation_spec,
                                             time_step_spec.reward)

compiler_opt/rl/policy_saver.py

@@ -181,7 +181,8 @@ def _write_output_signature(
 
     # Map spec name to index in flattened outputs.
     sm_action_indices = {
-        k.name.lower(): i for i, k in enumerate(sm_action_signature)}
+        k.name.lower(): i for i, k in enumerate(sm_action_signature)
+    }
 
     # List mapping flattened structured outputs to tensors.
     sm_action_tensors = saved_model.signatures['action'].outputs

compiler_opt/rl/regalloc/config.py

@@ -33,22 +33,26 @@ def get_regalloc_signature_spec():
 
   observation_spec = {
       key: tf.TensorSpec(dtype=tf.int64, shape=(num_registers), name=key)
-      for key in ('mask', 'is_hint', 'is_local', 'is_free')}
-  observation_spec.update(
-      {key: tensor_spec.BoundedTensorSpec(
-                dtype=tf.int64,
-                shape=(num_registers),
-                name=key,
-                minimum=0,
-                maximum=6) for key in ('max_stage', 'min_stage')})
-  observation_spec.update(
-      {key: tf.TensorSpec(dtype=tf.float32, shape=(num_registers), name=key)
-           for key in ('weighed_reads_by_max', 'weighed_writes_by_max',
-                       'weighed_read_writes_by_max', 'weighed_indvars_by_max',
-                       'hint_weights_by_max', 'start_bb_freq_by_max',
-                       'end_bb_freq_by_max', 'hottest_bb_freq_by_max',
-                       'liverange_size', 'use_def_density', 'nr_defs_and_uses',
-                       'nr_broken_hints', 'nr_urgent', 'nr_rematerializable')})
+      for key in ('mask', 'is_hint', 'is_local', 'is_free')
+  }
+  observation_spec.update({
+      key:
+          tensor_spec.BoundedTensorSpec(
+              dtype=tf.int64,
+              shape=(num_registers),
+              name=key,
+              minimum=0,
+              maximum=6) for key in ('max_stage', 'min_stage')
+  })
+  observation_spec.update({
+      key: tf.TensorSpec(dtype=tf.float32, shape=(num_registers), name=key)
+      for key in ('weighed_reads_by_max', 'weighed_writes_by_max',
+                  'weighed_read_writes_by_max', 'weighed_indvars_by_max',
+                  'hint_weights_by_max', 'start_bb_freq_by_max',
+                  'end_bb_freq_by_max', 'hottest_bb_freq_by_max',
+                  'liverange_size', 'use_def_density', 'nr_defs_and_uses',
+                  'nr_broken_hints', 'nr_urgent', 'nr_rematerializable')
+  })
   observation_spec['progress'] = tensor_spec.BoundedTensorSpec(
       dtype=tf.float32, shape=(), name='progress', minimum=0, maximum=1)