make jax.debug.print work with non-jax types

[dlwh]

For newbies/convenience, it'd be nice if jax.debug.print worked with static data, e.g. dtypes

import jax
import jax.numpy as jnp

@jax.jit
def example_function(input_array):
    jax.debug.print("input_array dtype: {dtype}", dtype=input_array.dtype)

# Example usage
input_array = jnp.array([1, 2, 3], dtype=jnp.float32)
example_function(input_array)

Currently we get:

TypeError: Value dtype('float32') with type <class 'numpy.dtype[float32]'> is not a valid JAX type

It'd be nice if the experience was the same as outside jit (where it works just fine), but a better error message would be ok too.

Please:

• Check for duplicate requests.
• Describe your goal, and if possible provide a code snippet with a motivating example.

[soraros]

You can use Python print for static I think?

[dlwh]

You can. It confused a newbie I was working with. The error message is not super suggestive of what's wrong.

You could also imagine wanting

 jax.debug.print("input_array {array}, dtype: {dtype}", array=input_array, dtype=input_array.dtype)

[dlwh]

(If you decide to just do better error message, I can upstream something to equinox since it's a very equinox-y thing to want)

[yashk2810]

No, I'll just make it work hopefully.