DOC: jax.lax.top_k: fix docstring rendering & add example#21788

copybara-service[bot] merged 1 commit intojax-ml:mainjax-ml/jax:mainfrom

jakevdp:top-k-docjakevdp/jax:top-k-docCopy head branch name to clipboard