Implement placeholder_value() for ScalarizerTraceType for tf.function…

… layering efforts. PiperOrigin-RevId: 497088543 Change-Id: I5aee396b6c97263a44ef5505cbb12a064dab5508
tensorflow · Dec 22, 2022 · 7071500 · 7071500
1 parent 7c1ef86
commit 7071500
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/tf_agents/bandits/multi_objective/multi_objective_scalarizer.py b/tf_agents/bandits/multi_objective/multi_objective_scalarizer.py
@@ -70,14 +70,17 @@ class ScalarizerTraceType(tf.types.experimental.TraceType):
   Protocol.
   """
 
+  def __init__(self, value):
+    self._value = value
+
   def is_subtype_of(self, _):
     return False
 
   def most_specific_common_supertype(self, _):
     return None
 
   def placeholder_value(self, placeholder_context=None):
-    raise NotImplementedError
+    return self._value
 
   def __hash__(self):
     return id(self)
@@ -182,7 +185,7 @@ def set_parameters(self, **kwargs):
 
   def __tf_tracing_type__(self, _):
     """Default TraceType Protocol for Scalarizaer Class."""
-    return ScalarizerTraceType()
+    return ScalarizerTraceType(self)
 
 
 class LinearScalarizer(Scalarizer):