[Doc] Update jnp.argmin/argmax doc for Pallas TPU#34623

catswe wants to merge 1 commit intojax-ml:mainjax-ml/jax:mainfrom

catswe:patch-1catswe/jax:patch-1Copy head branch name to clipboard