Martin Schubert

OK, I am currently working around this via, ```python def _eig_jax(matrix: jnp.ndarray) -> Tuple[jnp.ndarray, jnp.ndarray]: """Eigendecomposition using `jax.numpy.linalg.eig`.""" if jax.devices()[0] == jax.devices("cpu")[0]: return jnp.linalg.eig(matrix) else: eigvals, eigvecs = jax.pure_callback( _eig_jax_cpu,...

Thanks for your suggestion and I'll look forward to the new API. Unfortunately, I have found performance to be quite different (also when using `scipy.linalg.eig`). This is so both on...

Sure, here is some benchmarking code that I ran on colab CPU. I avoided using `%%timeit` (which reports mean time) since free tier colab seems to be quite noisy. ```python...

Hi @gdalle, nice work with `sparsediffax`. Your proposal is quite appealing---I don't have significant bandwidth to devote to this at the moment, but I am broadly interested in advancing the...

Hi @gdalle , thanks for this great summary. These all sound like excellent additions. To summarize, it seems like things fit into two categories: user-facing changes and those which are...

@gdalle This works for me! I am happy to have development take place in `sparsediffax`, and to point users toward that repo on the `sparsejax` homepage. And agreed that a...