commit d200783933197b35d487eeb554123a8f4f094e30
author Chase Naples <cnaples79@gmail.com> Tue Sep 23 00:29:53 2025 -0400
committer GitHub <noreply@github.com> Mon Sep 22 21:29:53 2025 -0700
tree fca8d14fb44d9e46183c9bd56dc1eb97db0584cf
parent cfcdd82a61fee7394a97b16cf8f579208650ccd8

[nfc] alias tf_agents.policies.policy_saver as tf_policy_saver (#528)

compiler_opt/rl/policy_saver.py

diff --git a/compiler_opt/rl/policy_saver.py b/compiler_opt/rl/policy_saver.py
index fc88e75..abf8aae 100644
--- a/compiler_opt/rl/policy_saver.py
+++ b/compiler_opt/rl/policy_saver.py
@@ -21,7 +21,7 @@
 
 import tensorflow as tf
 from tf_agents.policies import tf_policy
-from tf_agents.policies import policy_saver
+from tf_agents.policies import policy_saver as tf_policy_saver
 from tf_agents.typing import types as tf_agents_types
 
 OUTPUT_SIGNATURE = 'output_spec.json'
@@ -158,8 +158,8 @@
       policy_dict: A dict mapping from policy name to policy.
     """
     self._policy_saver_dict: dict[str, tuple[
-        policy_saver.PolicySaver, tf_policy.TFPolicy]] = {
-            policy_name: (policy_saver.PolicySaver(
+        tf_policy_saver.PolicySaver, tf_policy.TFPolicy]] = {
+            policy_name: (tf_policy_saver.PolicySaver(
                 policy, batch_size=1, use_nest_path_signatures=False), policy
                          ) for policy_name, policy in policy_dict.items()
         }