Handle bad_indices_policy for GatherNd and ScatterNd

PiperOrigin-RevId: 654697859
google-deepmind · Jul 22, 2024 · 2d2cb8d · 2d2cb8d
1 parent 12d8393
commit 2d2cb8d
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tf2jax/_src/ops.py b/tf2jax/_src/ops.py
@@ -2249,7 +2249,7 @@ def _func(xs: jnp.ndarray, shape: jnp.ndarray) -> ArrayLike:
 @register_operation("TensorScatterUpdate")
 def _tensor_scatter_update(proto):
   """Parse an TensorScatterUpdate Op."""
-  _check_attrs(proto, {"T", "Tindices"})
+  _check_attrs(proto, {"T", "Tindices", "bad_indices_policy"})
 
   def _func(
       operand: jnp.ndarray,