Implement "bad_indcies_policy" for ScatterNd.

For testing, we also introduced "ScatterNdTest" that verifies the existing "default" behavior for comparison. PiperOrigin-RevId: 637646019
google-deepmind · May 29, 2024 · 059bc5e · 059bc5e
1 parent ad9afbb
commit 059bc5e
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tf2jax/_src/ops.py b/tf2jax/_src/ops.py
@@ -1727,7 +1727,7 @@ def _func(
 @register_operation("ScatterNd")
 def _scatter_nd(proto):
   """Parse a ScatterNd op."""
-  _check_attrs(proto, {"T", "Tindices"})
+  _check_attrs(proto, {"T", "Tindices", "bad_indices_policy"})
 
   def _func(
       indices: jnp.ndarray,