jnp fails after enable x64 with Apple M1 chip

[willcai7]

Description

Right now I'm using jax with M1 pro chip. I cannot use functions from jnp after enable int64. This is my source code.

import jax
print(jax.__version__)
jax.config.update('jax_enable_x64', True)
print(jax.numpy.arange(10).dtype)
from jax import numpy as jnp
print(jnp.eye(3))

System info (python version, jaxlib version, accelerator, etc.)

print(jax.version)
0.4.26

[jakevdp]

Hi - thanks for the report. This is a known issue, tracked in #16435.

[willcai7]

Thank you!