Update expected attrs of VarHandleOp -- add debug_name.

PiperOrigin-RevId: 558314532
google-deepmind · Aug 19, 2023 · a66033d · a66033d
1 parent 22dfe41
commit a66033d
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/tf2jax/_src/ops.py b/tf2jax/_src/ops.py
@@ -2266,7 +2266,8 @@ def _func(x: jnp.ndarray) -> List[jnp.ndarray]:
 @register_operation("VarHandleOp")
 def _var_handle(proto):
   _check_attrs(
-      proto, {"shared_name", "container", "allowed_devices", "shape", "dtype"})
+      proto, {"shared_name", "container", "allowed_devices", "shape", "dtype",
+              "debug_name"})
 
   def _func():
     raise ValueError(f"VarHandleOp `{proto.name}` cannot be evaluated.")