Jax Eigen Value Decomposition for a NOT large matrix

[bethejulia]

Hi there

I am Using v2-8 TPU and have a code to diagonalize a symmetric Matrix.

import jax
import jax.numpy as jnp
from jax import random
N=5000 # the Dimension of Matrix
x = random.normal(random.PRNGKey(0), (N,N), dtype=jnp.float32)
y = jnp.matmul(x, x.T, precision=jax.lax.Precision.HIGHEST)
import datetime
now = datetime.datetime.now()
print (now.strftime("Time: %Y-%m-%d %H:%M:%S \n"))
from jax.lax.linalg import eigh as jeigh
f = jax.jit((jeigh))
f(y)
print (datetime.datetime.now().strftime("Time: %Y-%m-%d %H:%M:%S \n"))

It does Diagonalization of Symmetric Hermitian until when N=7000 (N is the size of Symm.Matrix)
But when N=8000 it crashed and give a Memory Error:

----------------------------------------------------------
jax._src.traceback_util.UnfilteredStackTrace: jaxlib.xla_extension.XlaRuntimeError: RESOURCE_EXHAUSTED: XLA:TPU compile permanent error. Ran out of memory in memory space hbm. Used 12.40G of 7.48G hbm. Exceeded hbm capacity by 4.92G.

Total hbm usage >= 12.92G:
    reserved        530.00M
    program          12.40G
    arguments            0B

Output size 0B; shares 0B with arguments.
----------------------------------------------------------

From the Documentation (https://jax.readthedocs.io/en/latest/changelog.html):
"The implementation of singular value decomposition (SVD) and symmetric/Hermitian eigendecomposition should be significantly faster on TPU, especially for matrices above 1000x1000 or so. Both now use a spectral divide-and-conquer algorithm for eigendecomposition (QDWH-eig)."

So Do I want to use v3-8 or greater TPU (like POD) to diagonalize a Matrix where N=8000 or beyond ?

Thanks in advance