[Python] Why does mlx.core.sqrt() crash on my MacBook Air M2 when applied to a complex...

mlx.core.sqrt() is crashing on my MacBook Air M2 when applied to a complex argument:

Python 3.11.11 (main, Dec 3 2024, 17:20:40) [Clang 16.0.0 (clang-1600.0.26.4)] on darwin
Type "help", "copyright", "credits" or "license" for more information.
>>> import mlx.core as mx
>>> x = mx.array([0.0, 1.0, 2.0])
>>> print(x)
array([0, 1, 2], dtype=float32)
>>> y = mx.sqrt(x)
>>> print(y)
array([0, 1, 1.41421], dtype=float32)
>>> w = 1j * x
>>> print(w)
array([0+0j, 0+1j, 0+2j], dtype=complex64)
>>> z = mx.sqrt(w)
>>> print(z)
Traceback (most recent call last):
File "<stdin>", line 1, in <module>
RuntimeError: [metal:evice] Unable to load kernel v_Sqrtcomplex64complex64

>>> print(mx.metal.device_info())
{'resource_limit': 499000, 'max_buffer_length': 8589934592, 'architecture': 'applegpu_g14g', 'memory_size': 17179869184, 'max_recommended_working_set_size': 11453251584, 'device_name': 'Apple M2'}


Is this expected?

I tried to debug it, using a capture, by running this executable Python script:

import mlx.core as np

GPU_TRACE_FNAME = "metal.gputrace"

metal_available = np.metal.is_available()
if metal_available:
print(f"Metal Device Info:\n{np.metal.device_info()}")
else:
print("Sorry, Metal is not available.")

np.metal.start_capture(GPU_TRACE_FNAME)
x = np.array([0.0, 1.0, 2.0])
y = 1j * x
z = np.sqrt(y)
print(z)
np.metal.stop_capture()


But, that fails:

% ./mlx_test_sqrt.py
Metal Device Info:
{'resource_limit': 499000, 'max_buffer_length': 8589934592, 'architecture': 'applegpu_g14g', 'memory_size': 17179869184, 'max_recommended_working_set_size': 11453251584, 'device_name': 'Apple M2'}
Traceback (most recent call last):
File ".../mlx_test_sqrt.py", line 22, in <module>
np.metal.start_capture(GPU_TRACE_FNAME)
RuntimeError: [metal::start_capture] Failed to start: Capturing is not supported.


Any idea what my next steps should be?

Continue reading...