Metal: Failed Assertion / Crash when trying to create arrays in double-precision

[johnjmolina]

Description

There seems to be an issue with casting for the edge case of size zero arrays.

from jax import config
config.update("jax_enable_x64", True)
import jax.numpy as jnp
import numpy as np

jnp.array([1.0])           # OK : Array([1.], dtype=float64)
jnp.array(np.double(1.0))  # OK : Array(1., dtype=float64)
jnp.array(jnp.double(1.0)) # Crash
jnp.array(1.0)             # Crash  

The last two commands will crash the kernel

Metal device set to: Apple M1

systemMemory: 16.00 GB
maxCacheSize: 5.33 GB

/AppleInternal/Library/BuildRoots/c2cb9645-dafc-11ed-aa26-6ec1e3b3f7b3/Library/Caches/com.apple.xbs/Sources/MetalPerformanceShadersGraph/mpsgraph/MetalPerformanceShadersGraph/Core/Files/MPSGraphExecutable.mm:1377: failed assertion `Incompatible element type for parameter at index 0, mlir module expected element type f64 but received f32'

What jax/jaxlib version are you using?

jax 0.4.11, jaxlib 0.4.10, jax-metal 0.0.2

Which accelerator(s) are you using?

Apple GPU

Additional system info

No response

NVIDIA GPU info

No response

[jakevdp]

Just to confirm, this is on the Apple metal backend, yes? (All of this should work fine on non-experimental backends)

[hawkinsp]

I think the Metal plugin simply doesn't support f64, but it should report it more gracefully.

[kulinseth]

@hawkinsp , that's right. Metal doesn't have FP64 support, we will add support to fail more gracefully.

[kulinseth]

@johnjmolina , this issue is not related to Empty shapes. We don't have support for FP64 and we will add error-checks early to exit gracefully.

[johnjmolina]

Thanks for the prompt reply, I understand the issue now.

[johannahaffner]

@hawkinsp , that's right. Metal doesn't have FP64 support, we will add support to fail more gracefully.

Are you planning on adding FP64 support? I would be really keen on that.

[shuhand0]

While we are looking into the issue, FP64 support will NOT be there for sometime.

[johannahaffner]

@shuhand0 thank you for letting me know.