sort feature of LAPACK's gees routine

[hyperfra]

I saw that the Schur implementation in jax.lax has as an input the ordering of the eigenvalues:

jax.lax.linalg.schur(x, *, compute_schur_vectors=True, sort_eig_vals=False, select_callable=None)

but the feature is not implemented : NotImplementedError: The sort feature of LAPACK's gees routine is not implemented.

Is there a way I could implement it? Or is the re-ordering not implemented for a deeper reason?

thanks

[jakevdp]

Hi - thanks for the question! The issue is that the semantics of sort in lapack are that it uses a callable select function to perform the sort (see https://netlib.org/lapack/explore-html/d9/d8e/group__double_g_eeigen_ga2233e8611b7b529422497036b284b438.html)

There is a TODO in the implementation to figure out how to pass such a callable function to the C++ routine: https://github.com/google/jax/blob/32026ad18bd767d2c1b72bae5fa1a65591defea3/jaxlib/lapack.py#L629

It's not entirely clear how to best do that, and so currently we default to sort='N', where the callable is not referenced.