[JAX] Adds the approx_top_k_p bridge.

PiperOrigin-RevId: 441172779

jax/experimental/jax2tf/jax2tf.py

@@ -983,7 +983,6 @@ def _unexpected_primitive(p: core.Primitive, *args, **kwargs):
     # Not high priority?
     "after_all",
     "all_to_all",
-    "approx_top_k",
     "create_token",
     "custom_transpose_call",
     "custom_vmap_call",
@@ -2380,6 +2379,23 @@ def promote_tf_dtype(tf_dtype):
 tf_impl[lax.top_k_p] = _top_k
 
 
+def _approx_top_k(operand: TfVal, k: int, reduction_dimension: int,
+                  recall_target: float, is_max_k: bool,
+                  reduction_input_size_override: int,
+                  aggregate_to_topk: bool) -> Tuple[TfVal, TfVal]:
+  if is_max_k:
+    return tf.math.approx_max_k(operand, k, reduction_dimension, recall_target,
+                                reduction_input_size_override,
+                                aggregate_to_topk)
+  else:
+    return tf.math.approx_min_k(operand, k, reduction_dimension, recall_target,
+                                reduction_input_size_override,
+                                aggregate_to_topk)
+
+
+tf_impl[lax.approx_top_k_p] = _approx_top_k
+
+
 def _sort(*operands: TfVal, dimension: int, is_stable: bool,
           num_keys: int) -> Tuple[TfVal, ...]:
   assert 1 <= num_keys <= len(operands)

jax/experimental/jax2tf/tests/jax2tf_limitations.py

@@ -181,6 +181,16 @@ def acosh(cls, harness: primitive_harness.Harness):
         cls.helper_get_trig_custom_limitation(np.cosh)
     ]
 
+  @classmethod
+  def approx_max_k(cls, harness: primitive_harness.Harness):
+    supported_dtypes = jtu.supported_dtypes()
+    return Jax2TfLimitation(
+        "eager is not supported in CPU or GPU.",
+        dtypes=[t for t in [jnp.bfloat16, np.float16, np.float32]
+                if t in supported_dtypes],
+        devices=("cpu", "gpu", "tpu"),
+        modes=("graph", "compiled"))
+
   @classmethod
   def argmax(cls, harness: primitive_harness.Harness):
     return [