Can't instantiate jax numpy zeros if 64-bit mode is enabled using GPU

[nelsen4]

Hi, im quite new to using GPUs. I just wanted to run a module implemented in jax, on a mac m2 GPU (jax==0.4.11, jax-metal==0.0.4, jaxlib==0.4.11). The code runs on CPU, but when using the GPU, i cant instantiate a jnp.zeros/ones array if 64 bit mode is set to true. Minimum example is:

import jax.numpy as jnp

from jax.config import config
config.update("jax_enable_x64", True)

results = jnp.ones((200, 30))

which throws:

jaxlib.xla_extension.XlaRuntimeError: UNKNOWN: <unknown>:0: error: 'func.func' op One or more function input/output data types are not supported.
<unknown>:0: note: see current operation: 
"func.func"() <{arg_attrs = [{mhlo.sharding = "{replicated}"}], function_type = (tensor<f64>) -> tensor<200x30xf64>, res_attrs = [{}], sym_name = "main", sym_visibility = "public"}> ({
^bb0(%arg0: tensor<f64>):
  %0 = "mhlo.broadcast_in_dim"(%arg0) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f64>) -> tensor<200x30xf64>
  "func.return"(%0) : (tensor<200x30xf64>) -> ()
}) : () -> ()
<unknown>:0: error: failed to legalize operation 'func.func'
<unknown>:0: note: see current operation: 
"func.func"() <{arg_attrs = [{mhlo.sharding = "{replicated}"}], function_type = (tensor<f64>) -> tensor<200x30xf64>, res_attrs = [{}], sym_name = "main", sym_visibility = "public"}> ({
^bb0(%arg0: tensor<f64>):
  %0 = "mhlo.broadcast_in_dim"(%arg0) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f64>) -> tensor<200x30xf64>
  "func.return"(%0) : (tensor<200x30xf64>) -> ()
}) : () -> ()

Not sure if im making some mistake or this is a known issue.
If the update flag is set to false, the code runs fine. And for jnp.array() the code also runs fine. Is this expected? (I guess zeros and ones wont be 64 bit anyways, but i dont understand why it would throw an error on gpu but not on cpu, maybe it has something to do with deprecating the 64 bit flag?)

Also later in the code, when trying to update values in the array, i get another issue which is known (https://developer.apple.com/forums/thread/738697).

[jakevdp]

We have GPU CI tests that cover things like 64-bit array creation. Can you say more about the architecture you are running on?

[nelsen4]

Sorry, forgot to add. Im running Sonoma 14.0 on an Apple M2 Pro (32GB RAM)

[jakevdp]

So is this using the Metal plugin?

[jakevdp]

Oh, I see you mention jax-metal above.

jax-metal is highly experimental, and does not support the full JAX API. You can see the current list of reported metal-related issues here: https://github.com/google/jax/issues?q=is%3Aopen+is%3Aissue+label%3A%22Apple+GPU+%28Metal%29+plugin%22

It sounds like you're running into the issue previously reported in #16435

Note that jax-metal is closed source and the JAX core team has no access to the code, so there's not much we can do about issues like this.

[nelsen4]

Thanks!