[Bug] Crash in `item()`.

[tqtifnypmb]

Code:

let tokenId: Int = logits[idx].item()

Occurs occasionally, not 100% reproducible

Crash Stacks:

0   libsystem_kernel.dylib        	0x00000001e9a86974 __pthread_kill + 8 (:-1)
1   libsystem_pthread.dylib       	0x00000001fd5090ec pthread_kill + 268 (pthread.c:1717)
2   libsystem_c.dylib             	0x00000001a94d7b80 abort + 180 (abort.c:118)
3   libc++abi.dylib               	0x00000001fd428ccc abort_message + 132 (abort_message.cpp:78)
4   libc++abi.dylib               	0x00000001fd418e68 demangling_terminate_handler() + 320 (cxa_default_handlers.cpp:72)
5   libobjc.A.dylib               	0x00000001994860cc _objc_terminate() + 160 (objc-exception.mm:499)
6   libc++abi.dylib               	0x00000001fd428090 std::__terminate(void (*)()) + 16 (cxa_handlers.cpp:59)
7   libc++abi.dylib               	0x00000001fd42b2e4 __cxxabiv1::failed_throw(__cxxabiv1::__cxa_exception*) + 88 (cxa_exception.cpp:152)
8   libc++abi.dylib               	0x00000001fd42b228 __cxa_throw + 308 (cxa_exception.cpp:283)
9  Demo              	                 0x000000010360acf0 mlx::core::reshape(mlx::core::array const&, std::__1::vector<int, std::__1::allocator<int>>, std::__1::variant<std::__1::monostate, mlx::core::Stream, mlx::core::Device>) + 916 (ops.cpp:328)
10  Demo              	                 0x00000001031cb78c mlx_reshape + 108 (ops.cpp:633)
11  Demo              	                 0x00000001036c2cd8 MLXArray.reshaped(_:stream:) + 256 (MLXArray+Ops.swift:1974)
12  Demo              	                 0x00000001036b2610 getItemND(src:operations:stream:) + 2404 (MLXArray+Indexing.swift:722)
13  Demo              	                 0x00000001036b5190 MLXArray.subscript.getter + 436

Any hints about what is happening behind the scenes or how to fix it?

Thanks