Hello, I'm interested in trying the new JAX Metal plug-in and followed the steps in https://developer.apple.com/metal/jax/. Upon installation, I don't see any difference between the backend device detected by JAX and a pure CPU setup: >>> import jax >>> jax.devices() [CpuDevice(id=0)] >>> jax.devices()[0].platform 'cpu' >>> jax.devices()[0].device_kind 'cpu' >>> jax.devices()[0].client.platform 'cpu' >>> jax.devices()[0].client.runtime_type 'tfrt' Is this really using a Metal backend? How can I determine for sure? Thank you!

Hello, The metal compiler is crashing for me when attempting to compile a metal source file that contains Core Image kernel implementations. This is a minimal version of a file that produces the crash: extern "C" { namespace coreimage { inline void swap(thread float4 &a, thread float4 &b) { float4 tmp = a; a = min(a, b); b = max(tmp, b); } typedef sample_t s; float4 median_reduction_3(s v0, s v1, s v2) { swap(v1, v2); swap(v0, v2); swap(v0, v1); return v1; } }} Some observations: If inline is removed, the code compiles fine. I'm not sure if there's a performance impact, as the backend llvm compiler might as well decide to inline it on its own. If the calls to swap are commented in the median reduction function, the code compiles. If the -fcikernel compilation flag is not used, it also compiles fine (doesn't crash). Of course, that configuration doesn't allow the use of functions inside the file as Core Image kernels. I'm using the build settings recommended in this WWDC20 session (without indicating the location of the header files, since it's empty in my project and the new compiler interprets the argument following -I as a directory).