Matrixexponential: "expm(0) == nan"

I have noticed that feeding an array of only zeros into jax.scipy.linalg.expm results in a purely np.nan array. It should be the unit matrix, e.g. expm(zeros) == np.eye(*zeros.shape)

Minimal reproducer:

import jax.numpy as np
from jax.version import __version__
from jax.scipy.linalg import expm

print(f'jax version: {__version__}')
#  jax version: 0.1.58

zero = np.zeros([2, 2], dtype=np.float32)
exp_zero = expm(zero)
print(f'exp(zero) == \n{exp_zero}')
#  exp(zero) == 
#  [[nan nan]
#  [nan nan]]

[shoyer]

• @jhueckelheim @sriharikrishna

[shoyer]

Should be fixed now by #2192