Add path prefix to policy utils test (#372)

This patch adds a path prefix to all the uses of vocab within the policy utils test so that a path prefix can easily be added in downstream repository instances.
google · Sep 17, 2024 · 00f9e4a · 00f9e4a
1 parent 5230947
commit 00f9e4a
Showing 1 changed file with 12 additions and 5 deletions.
diff --git a/compiler_opt/es/policy_utils_test.py b/compiler_opt/es/policy_utils_test.py
@@ -32,6 +32,8 @@
 from compiler_opt.rl.regalloc import config as regalloc_config
 from compiler_opt.rl.regalloc import regalloc_network
 
+TEST_PATH_PREFIX = ''
+
 
 class ConfigTest(absltest.TestCase):
 
@@ -40,7 +42,8 @@ def test_inlining_config(self):
     problem_config = registry.get_configuration(
         implementation=inlining.InliningConfig)
     time_step_spec, action_spec = problem_config.get_signature_spec()
-    quantile_file_dir = os.path.join('compiler_opt', 'rl', 'inlining', 'vocab')
+    quantile_file_dir = os.path.join(TEST_PATH_PREFIX, 'compiler_opt', 'rl',
+                                     'inlining', 'vocab')
     creator = inlining_config.get_observation_processing_layer_creator(
         quantile_file_dir=quantile_file_dir,
         with_sqrt=False,
@@ -71,7 +74,8 @@ def test_regalloc_config(self):
     problem_config = registry.get_configuration(
         implementation=regalloc.RegallocEvictionConfig)
     time_step_spec, action_spec = problem_config.get_signature_spec()
-    quantile_file_dir = os.path.join('compiler_opt', 'rl', 'regalloc', 'vocab')
+    quantile_file_dir = os.path.join(TEST_PATH_PREFIX, 'compiler_opt', 'rl',
+                                     'regalloc', 'vocab')
     creator = regalloc_config.get_observation_processing_layer_creator(
         quantile_file_dir=quantile_file_dir,
         with_sqrt=False,
@@ -113,7 +117,8 @@ def test_set_vectorized_parameters_for_policy(self):
     problem_config = registry.get_configuration(
         implementation=inlining.InliningConfig)
     time_step_spec, action_spec = problem_config.get_signature_spec()
-    quantile_file_dir = os.path.join('compiler_opt', 'rl', 'inlining', 'vocab')
+    quantile_file_dir = os.path.join(TEST_PATH_PREFIX, 'compiler_opt', 'rl',
+                                     'inlining', 'vocab')
     creator = inlining_config.get_observation_processing_layer_creator(
         quantile_file_dir=quantile_file_dir,
         with_sqrt=False,
@@ -176,7 +181,8 @@ def test_get_vectorized_parameters_from_policy(self):
     problem_config = registry.get_configuration(
         implementation=inlining.InliningConfig)
     time_step_spec, action_spec = problem_config.get_signature_spec()
-    quantile_file_dir = os.path.join('compiler_opt', 'rl', 'inlining', 'vocab')
+    quantile_file_dir = os.path.join(TEST_PATH_PREFIX, 'compiler_opt', 'rl',
+                                     'inlining', 'vocab')
     creator = inlining_config.get_observation_processing_layer_creator(
         quantile_file_dir=quantile_file_dir,
         with_sqrt=False,
@@ -224,7 +230,8 @@ def test_tfpolicy_and_loaded_policy_produce_same_variable_order(self):
     problem_config = registry.get_configuration(
         implementation=inlining.InliningConfig)
     time_step_spec, action_spec = problem_config.get_signature_spec()
-    quantile_file_dir = os.path.join('compiler_opt', 'rl', 'inlining', 'vocab')
+    quantile_file_dir = os.path.join(TEST_PATH_PREFIX, 'compiler_opt', 'rl',
+                                     'inlining', 'vocab')
     creator = inlining_config.get_observation_processing_layer_creator(
         quantile_file_dir=quantile_file_dir,
         with_sqrt=False,