$ python 
Python 3.12.2 (main, Feb  6 2024, 20:19:44) [Clang 15.0.0 (clang-1500.1.0.2.5)] on darwin
Type "help", "copyright", "credits" or "license" for more information.
>>> import jax.numpy as jnp 
>>> arr1 = jnp.array([[[2, 5, 0]]], dtype=jnp.uint8) # default datatype int32 actually does work 
Platform 'METAL' is experimental and not all JAX functionality may be correctly supported!
2024-03-24 16:15:50.307903: W pjrt_plugin/src/mps_client.cc:563] WARNING: JAX Apple GPU support is experimental and not all JAX functionality is correctly supported!
Metal device set to: Apple M1 Pro

systemMemory: 32.00 GB
maxCacheSize: 10.67 GB

>>> arr2 = jnp.array([12, 13,  1]) # both uint8 and default work 
>>> jnp.tile(arr1, arr2)
loc("jit(reshape)/jit(main)/reshape[new_sizes=(1, 1, 1, 1, 1, 3) dimensions=None]"("<stdin>":1:0)): error: 'anec.reshape' op result #0 must be 4D/5D memref of 16-bit float or 8-bit signed integer or 8-bit unsigned integer values, but got 'memref<1x1x1x1x1x3xui8>'
zsh: abort      python

 $ python                
Python 3.12.2 (main, Feb  6 2024, 20:19:44) [Clang 15.0.0 (clang-1500.1.0.2.5)] on darwin
Type "help", "copyright", "credits" or "license" for more information.
>>> import jax.numpy as jnp 
r>>> 
>>> arr = jnp.arange(20)
Platform 'METAL' is experimental and not all JAX functionality may be correctly supported!
2024-03-24 16:21:53.255063: W pjrt_plugin/src/mps_client.cc:563] WARNING: JAX Apple GPU support is experimental and not all JAX functionality is correctly supported!
Metal device set to: Apple M1 Pro

systemMemory: 32.00 GB
maxCacheSize: 10.67 GB

>>> arr = jnp.arange(20).reshape((1,1,1,1,1,5,4)) # does not error? 
>>> arr = jnp.arange(20, dtype=jnp.uint8).reshape((1,1,1,1,1,5,4)) # has the error but does not abort the kernel
loc("jit(reshape)/jit(main)/reshape[new_sizes=(1, 1, 1, 1, 1, 5, 4) dimensions=None]"("<stdin>":1:6)): error: 'anec.reshape' op result #0 must be 4D/5D memref of 16-bit float or 8-bit signed integer or 8-bit unsigned integer values, but got 'memref<1x1x1x1x1x5x4xui8>'
>>> arr
Array([[[[[[[ 0,  1,  2,  3],
            [ 4,  5,  6,  7],
            [ 8,  9, 10, 11],
            [12, 13, 14, 15],
            [16, 17, 18, 19]]]]]]], dtype=uint8)
>>> arr = jnp.arange(20, dtype=jnp.uint8).reshape((1,1,1,1,1,5,4)) # unsure why this doesn't error, assuming some under-the-hood behavior
>>> arr = jnp.arange(21, dtype=jnp.uint8).reshape((1,1,1,1,1,3,7))
loc("jit(reshape)/jit(main)/reshape[new_sizes=(1, 1, 1, 1, 1, 3, 7) dimensions=None]"("<stdin>":1:6)): error: 'anec.reshape' op result #0 must be 4D/5D memref of 16-bit float or 8-bit signed integer or 8-bit unsigned integer values, but got 'memref<1x1x1x1x1x3x7xui8>'
>>> 

 $ pip freeze | grep jax 
jax==0.4.20
jax-metal==0.0.5
jaxlib==0.4.20