From 9c81ac67cf73932051993103efeae1e0617b7aa9 Mon Sep 17 00:00:00 2001 From: Aiden Grossman Date: Mon, 16 Sep 2024 19:07:34 -0700 Subject: [PATCH] Fix some import related style guide violations in ES (#371) This patch fixes some import related style guide violations, particularly the importing of multiple modules/packages on a single line. Some internal tooling is not able to effectively analyze these patterns (given they don't show up often since they're banned by the style guide) which is relatively annoying. We also should not be importing individual classes/functions, so fix those too when they coincide with the previous point. --- compiler_opt/es/blackbox_learner_test.py | 8 ++- compiler_opt/es/blackbox_optimizers_test.py | 73 ++++++++++++--------- compiler_opt/es/es_trainer.py | 4 +- compiler_opt/es/es_trainer_lib.py | 3 +- 4 files changed, 52 insertions(+), 36 deletions(-) diff --git a/compiler_opt/es/blackbox_learner_test.py b/compiler_opt/es/blackbox_learner_test.py index 5f74a13a..348e90bd 100644 --- a/compiler_opt/es/blackbox_learner_test.py +++ b/compiler_opt/es/blackbox_learner_test.py @@ -27,9 +27,13 @@ from compiler_opt.distributed import worker from compiler_opt.distributed.local import local_worker_manager -from compiler_opt.es import blackbox_learner, policy_utils +from compiler_opt.es import blackbox_learner +from compiler_opt.es import policy_utils from compiler_opt.es import blackbox_optimizers -from compiler_opt.rl import corpus, inlining, policy_saver, registry +from compiler_opt.rl import corpus +from compiler_opt.rl import inlining +from compiler_opt.rl import policy_saver +from compiler_opt.rl import registry from compiler_opt.rl.inlining import config as inlining_config from compiler_opt.es import blackbox_evaluator diff --git a/compiler_opt/es/blackbox_optimizers_test.py b/compiler_opt/es/blackbox_optimizers_test.py index 2de73a57..16dc086b 100644 --- a/compiler_opt/es/blackbox_optimizers_test.py +++ b/compiler_opt/es/blackbox_optimizers_test.py @@ -52,7 +52,6 @@ import numpy as np from compiler_opt.es import blackbox_optimizers -from compiler_opt.es.blackbox_optimizers import EstimatorType, UpdateMethod, RegressionType from compiler_opt.es import gradient_ascent_optimization_algorithms perturbation_array = np.array([[0, 1], [2, -1], [4, 2], @@ -65,10 +64,12 @@ class BlackboxOptimizationAlgorithmsTest(parameterized.TestCase): @parameterized.parameters( - (perturbation_array, function_value_array, EstimatorType.ANTITHETIC, 3, + (perturbation_array, function_value_array, + blackbox_optimizers.EstimatorType.ANTITHETIC, 3, np.array([[4, 2], [2, 6], [-1, 5], [-2, -2], [8, -6], [1, -5] ]), np.array([10, -8, 4, -10, 8, -4])), - (perturbation_array, function_value_array, EstimatorType.FORWARD_FD, 5, + (perturbation_array, function_value_array, + blackbox_optimizers.EstimatorType.FORWARD_FD, 5, np.array([[4, 2], [8, -6], [-1, 5], [0, -3], [2, -1] ]), np.array([10, 8, 4, 2, 1]))) def test_filtering(self, perturbations, function_values, est_type, @@ -79,21 +80,23 @@ def test_filtering(self, perturbations, function_values, est_type, np.testing.assert_array_equal(expected_fs, top_fs) @parameterized.parameters( - (perturbation_array, function_value_array, EstimatorType.ANTITHETIC, 3, - np.array([100, -16])), (perturbation_array, function_value_array, - EstimatorType.FORWARD_FD, 5, np.array([76, -9])), - (perturbation_array, function_value_array, EstimatorType.ANTITHETIC, 0, - np.array([102, -34])), - (perturbation_array, function_value_array, EstimatorType.FORWARD_FD, 0, - np.array([74, -34]))) + (perturbation_array, function_value_array, + blackbox_optimizers.EstimatorType.ANTITHETIC, 3, np.array([100, -16])), + (perturbation_array, function_value_array, + blackbox_optimizers.EstimatorType.FORWARD_FD, 5, np.array([76, -9])), + (perturbation_array, function_value_array, + blackbox_optimizers.EstimatorType.ANTITHETIC, 0, np.array([102, -34])), + (perturbation_array, function_value_array, + blackbox_optimizers.EstimatorType.FORWARD_FD, 0, np.array([74, -34]))) def test_monte_carlo_gradient(self, perturbations, function_values, est_type, num_top_directions, expected_gradient): precision_parameter = 0.1 step_size = 0.01 current_value = 2 blackbox_object = blackbox_optimizers.MonteCarloBlackboxOptimizer( - precision_parameter, est_type, False, UpdateMethod.NO_METHOD, None, - step_size, num_top_directions) + precision_parameter, est_type, False, + blackbox_optimizers.UpdateMethod.NO_METHOD, None, step_size, + num_top_directions) current_input = np.zeros(2) step = blackbox_object.run_step(perturbations, function_values, current_input, current_value) @@ -106,13 +109,14 @@ def test_monte_carlo_gradient(self, perturbations, function_values, est_type, np.testing.assert_array_almost_equal(expected_gradient, gradient) @parameterized.parameters( - (perturbation_array, function_value_array, EstimatorType.ANTITHETIC, 3, - np.array([100, -16])), (perturbation_array, function_value_array, - EstimatorType.FORWARD_FD, 5, np.array([76, -9])), - (perturbation_array, function_value_array, EstimatorType.ANTITHETIC, 0, - np.array([102, -34])), - (perturbation_array, function_value_array, EstimatorType.FORWARD_FD, 0, - np.array([74, -34]))) + (perturbation_array, function_value_array, + blackbox_optimizers.EstimatorType.ANTITHETIC, 3, np.array([100, -16])), + (perturbation_array, function_value_array, + blackbox_optimizers.EstimatorType.FORWARD_FD, 5, np.array([76, -9])), + (perturbation_array, function_value_array, + blackbox_optimizers.EstimatorType.ANTITHETIC, 0, np.array([102, -34])), + (perturbation_array, function_value_array, + blackbox_optimizers.EstimatorType.FORWARD_FD, 0, np.array([74, -34]))) def test_monte_carlo_gradient_with_gradient_ascent_optimizer( self, perturbations, function_values, est_type, num_top_directions, expected_gradient): @@ -124,8 +128,9 @@ def test_monte_carlo_gradient_with_gradient_ascent_optimizer( step_size, 0.0)) blackbox_object = ( blackbox_optimizers.MonteCarloBlackboxOptimizer( - precision_parameter, est_type, False, UpdateMethod.NO_METHOD, None, - None, num_top_directions, gradient_ascent_optimizer)) + precision_parameter, est_type, False, + blackbox_optimizers.UpdateMethod.NO_METHOD, None, None, + num_top_directions, gradient_ascent_optimizer)) current_input = np.zeros(2) step = blackbox_object.run_step(perturbations, function_values, current_input, current_value) @@ -137,15 +142,18 @@ def test_monte_carlo_gradient_with_gradient_ascent_optimizer( np.testing.assert_array_almost_equal(expected_gradient, gradient) - @parameterized.parameters( - (perturbation_array, function_value_array, EstimatorType.ANTITHETIC, 3, - np.array([0.00483, -0.007534])), - (perturbation_array, function_value_array, EstimatorType.FORWARD_FD, 5, - np.array([0.012585, 0.000748])), - (perturbation_array, function_value_array, EstimatorType.ANTITHETIC, 0, - np.array([0.019319, -0.030134])), - (perturbation_array, function_value_array, EstimatorType.FORWARD_FD, 0, - np.array([0.030203, 0.001796]))) + @parameterized.parameters((perturbation_array, function_value_array, + blackbox_optimizers.EstimatorType.ANTITHETIC, 3, + np.array([0.00483, -0.007534])), + (perturbation_array, function_value_array, + blackbox_optimizers.EstimatorType.FORWARD_FD, 5, + np.array([0.012585, 0.000748])), + (perturbation_array, function_value_array, + blackbox_optimizers.EstimatorType.ANTITHETIC, 0, + np.array([0.019319, -0.030134])), + (perturbation_array, function_value_array, + blackbox_optimizers.EstimatorType.FORWARD_FD, 0, + np.array([0.030203, 0.001796]))) def test_sklearn_gradient(self, perturbations, function_values, est_type, num_top_directions, expected_gradient): precision_parameter = 0.1 @@ -156,8 +164,9 @@ def test_sklearn_gradient(self, perturbations, function_values, est_type, gradient_ascent_optimization_algorithms.MomentumOptimizer( step_size, 0.0)) blackbox_object = blackbox_optimizers.SklearnRegressionBlackboxOptimizer( - RegressionType.RIDGE, regularizer, est_type, True, - UpdateMethod.NO_METHOD, [], None, gradient_ascent_optimizer) + blackbox_optimizers.RegressionType.RIDGE, regularizer, est_type, True, + blackbox_optimizers.UpdateMethod.NO_METHOD, [], None, + gradient_ascent_optimizer) current_input = np.zeros(2) step = blackbox_object.run_step(perturbations, function_values, current_input, current_value) diff --git a/compiler_opt/es/es_trainer.py b/compiler_opt/es/es_trainer.py index c0e43807..f424b5c6 100644 --- a/compiler_opt/es/es_trainer.py +++ b/compiler_opt/es/es_trainer.py @@ -14,7 +14,9 @@ # limitations under the License. """Local ES trainer.""" -from absl import app, flags, logging +from absl import app +from absl import flags +from absl import logging import gin from compiler_opt.es import es_trainer_lib diff --git a/compiler_opt/es/es_trainer_lib.py b/compiler_opt/es/es_trainer_lib.py index fd4eb770..21a220bf 100644 --- a/compiler_opt/es/es_trainer_lib.py +++ b/compiler_opt/es/es_trainer_lib.py @@ -25,7 +25,8 @@ from compiler_opt.es import gradient_ascent_optimization_algorithms from compiler_opt.es import blackbox_learner from compiler_opt.es import policy_utils -from compiler_opt.rl import policy_saver, corpus +from compiler_opt.rl import policy_saver +from compiler_opt.rl import corpus POLICY_NAME = "policy"