Kernel always restarting when JIT compiling the forward call on MacBook Pro M3 Max

[Gdesau]

Hello,

I'm trying to use whisper-jax on a MacBook Pro M3 Max (with CPU 16 cores, GPU 40 cores, Neural Engine 16 cores and 128Gb unified memory). I'm working on a Jupyter Notebook and I followed the Apple Dev Doc to install JAX. I just had to add one more step to the doc: pip install jax==0.4.26 jaxlib==0.4.26 because I ran into an error(see the Apple Dev forum. However every time I try to JIT compile the forward call as in the doc, my kernel is restarting. Here is my code:

# Imports
from whisper_jax import FlaxWhisperPipline
import jax.numpy as jnp
from jax.experimental.compilation_cache import compilation_cache as cc
from IPython.display import Audio

# instantiate pipeline
pipeline = FlaxWhisperPipline("openai/whisper-small", dtype=jnp.bfloat16, batch_size=16)

cc.set_cache_dir("./jax_cache")

audio = './test_court.mp3'
Audio(audio)

%%time
# JIT compile the forward call - slow, but we only do once
text = pipeline(audio)

I've already tried with other versions of the model and I've tried in a .py file on VsCode and no error but nothing happened when printing the text.

Any idea?

Thanks in advance for the help