TopK
====

Retrieve the top-K largest or smallest elements along a specified axis. Given an input tensor of shape [a_0, a_1, ..., a_{n-1}] and integer argument k, return two outputs: * Value tensor of shape [a_0, a_1, ..., a_{axis-1}, k, a_{axis+1}, ... a_{n-1}] which contains the values of the top k elements ...

**Abstract Signature:**

``TopK(X: Tensor, K: int, axis: int, largest: int, sorted: int)``

.. raw:: html

 
    <div class="op-tabs-container"> 
      <div class="op-tabs-nav"> 
        <button class="op-tab-btn active" onclick="openOpTab(event, 'PyTorch_0')">PyTorch</button><button class="op-tab-btn " onclick="openOpTab(event, 'JAX (Core)_1')">JAX (Core)</button><button class="op-tab-btn " onclick="openOpTab(event, 'Apple MLX_2')">Apple MLX</button> 
      </div> 
      <div class="op-tabs-content"> 
        <div id="PyTorch_0" class="op-tab-pane active"><h4>PyTorch</h4><div class="op-detail-row"><span class="label">API:</span> <code>torch.topk</code></div><div class="op-detail-row"><span class="label">Strategy:</span> <span>Direct Mapping</span></div><div class="op-detail-row"><a href="https://pytorch.org/docs/stable/generated/torch.topk.html" target="_blank" class="op-doc-link">Official Docs ↗</a></div></div><div id="JAX (Core)_1" class="op-tab-pane "><h4>JAX (Core)</h4><div class="op-detail-row"><span class="label">API:</span> <code>jax.lax.top_k</code></div><div class="op-detail-row"><span class="label">Strategy:</span> <span>Plugin (topk_adapter)</span></div></div><div id="Apple MLX_2" class="op-tab-pane "><h4>Apple MLX</h4><div class="op-detail-row"><span class="label">API:</span> <code>mlx.core.topk</code></div><div class="op-detail-row"><span class="label">Strategy:</span> <span>Direct Mapping</span></div><div class="op-detail-row"><a href="https://ml-explore.github.io/mlx/build/html/python/_autosummary/mlx.core.topk.html" target="_blank" class="op-doc-link">Official Docs ↗</a></div></div> 
      </div> 

      <!-- Load JS Logic only once per page ideally, but safe to exist globally --> 
      <script src="../_static/op_tabs.js"></script> 
    </div>