NotImplementedError: MLIR translation rule for primitive 'debug_callback' not found for platform METAL

[amavrits]

Description

I am running jax with metal and encounter the following issue (along with a much longer tracebck):

File ".../.venv/lib/python3.11/site-packages/jax/_src/interpreters/mlir.py", line 2041, in lower_per_platform
    raise NotImplementedError(
NotImplementedError: MLIR translation rule for primitive 'debug_callback' not found for platform METAL

It appears that jax-tqdm is the culprit. Upon disabling it, the issue is resolved.

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

Python 3.11.9
jax 0.5.0
macOS Sequoia 15.5

[shuhand0]

Can you provide a simplified script for me to reproduce? metal-plugin probably has no support for debug_callback. We will look into it if it is a mandatory feature for plugin platform, especially from PJRT C API.

[ZagButNoZig]

import jax

@jax.jit
def foo():
    jax.debug.print("{}", 4)

foo()

This reproduces it for me @shuhand0

[amavrits]

@shuhand0 Apologies for the late reply. Indeed, the snippet by @ZagButNoZig reproduces the error.