Improve docstring for optax.power_iteration.

[carlosgmartin]

Minor improvements to the docstring for optax.power_iteration, including explaining the connection to the spectral radius of the matrix and noting that a dominant eigenvalue might not exist.

[vroulet]

Thanks!

[vroulet]

Sorry thinking a bit more about it, the wikipedia reference should be enough for additional mathematical details. I'm not sure we need to lengthen the docs here.

[carlosgmartin]

The reason I added this is because the current doc makes it seem like a ("the") dominant eigenvalue exists. But it might not, and in that case, the algorithm fails to converge. Also, one of the common uses of this function is computing the spectral radius, but the function might not always work for that end, for the aforementioned reason.

That's the intent behind my clarification.

[vroulet]

@carlosgmartin, thanks for the clarification.
Would the following note be good with you (I'll take care of the merge internally)?

  .. note::
    If the matrix is not diagonalizable or the dominant eigenvalue is not unique, the algorithm may not converge. 

and add at the beginning of the docstring

dominant eigenvalue (i.e, the spectral radius)

[carlosgmartin]

Sure.

[carlosgmartin]

@vroulet To clarify, should I do anything, or do you take it from here?

[vroulet]

I took care of it, it should get merged soon. Thanks again