From f937567d800baa3e62c1e5501b1687ab71907c7e Mon Sep 17 00:00:00 2001
From: Aiden Grossman <aidengrossman@google.com>
Date: Tue, 17 Sep 2024 00:23:25 +0000
Subject: [PATCH] Add path prefix to policy utils test

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.
---
 compiler_opt/es/policy_utils_test.py | 17 ++++++++++++-----
 1 file changed, 12 insertions(+), 5 deletions(-)

diff --git a/compiler_opt/es/policy_utils_test.py b/compiler_opt/es/policy_utils_test.py
index b266ceb3..e94a67db 100644
--- 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,