jax.numpy.diag\_indices ======================= .. currentmodule:: jax.numpy .. autofunction:: diag_indices