ApproxMaxK

Returns max k values and their indices using approximate search (TPU optimized).

Abstract Signature:

ApproxMaxK(operand: Tensor, k: int, reduction_dimension: int = -1)

PyTorch

API: torch.topk
Strategy: Direct Mapping
Official Docs ↗

JAX (Core)

API: jax.lax.approx_max_k
Strategy: Direct Mapping
Official Docs ↗

TensorFlow

API: tf.math.approx_max_k
Strategy: Direct Mapping
Official Docs ↗

Apple MLX

API: mlx.core.topk
Strategy: Direct Mapping
Official Docs ↗