Eigen-decomposition jitted on CPU in TPU backend process.

[yonghakim]

Hi,

This is similar question to 'Fall back to CPU for jax.scipy.linalg.schur?' #12245
and related to 'pure_callback for eigendecomposition' #14199

I am trying to use GPU or TPU for matrix operation except eigen-decomposition which is jitted on CPU since it's not implemented.

Below is the code I used to do it on CPU. It worked in GPU backend but failed in TPU backend with following error NotImplementedError: 64-bit types not supported..

import jax
import jax.numpy as jnp

def eig(X, type_complex):
    import numpy as np
    jax.config.update('jax_enable_x64', True)

    type_complex = jnp.complex64
    X = jnp.array(X, dtype=jnp.complex64)
    _eig = lambda x: jax.jit(np.linalg.eig, device=jax.devices('cpu')[0])(x)

    eigenvalues_shape = jax.ShapeDtypeStruct(X.shape[:-1], type_complex)
    eigenvectors_shape = jax.ShapeDtypeStruct(X.shape, type_complex)

    result_shape_dtype = (eigenvalues_shape, eigenvectors_shape)
    return jax.pure_callback(_eig, result_shape_dtype, X)

X = jnp.arange(1000*1000).reshape((1000,1000))
eig(X, jnp.complex64)

Is there a workaround that I can enforce eigen-decomposition on CPU while others run on TPU?

Thanks for reading.