keras/src/backend/jax/numpy.py/diagonal
def diagonal(x, offset=0, axis1=0, axis2=1):
    x = convert_to_tensor(x)
    return jnp.diagonal(
        x,
        offset=offset,
        axis1=axis1,
        axis2=axis2,
    )