Solving nonlinear problem on Metal GPU crashes Julia session

[ErikQQY]

Describe the bug 🐞

The testing example works fine with CUDA, but it does not work with Metal.jl and crashes the Julia session. But SimpleNonlinearSolve.jl works fine.

Expected behavior

A clear and concise description of what you expected to happen.

Minimal Reproducible Example 👇

Without MRE, we would only be able to help you to a limited extent, and attention to the issue would be limited. to know more about MRE refer to wikipedia and stackoverflow.

using Metal, NonlinearSolve
f(u, p) = u .* u .- 2
u0 = MtlArray([1.0f0, 1.0f0])
prob = NonlinearProblem(f, u0)
solver = solve(prob, NewtonRaphson())

Error & Stacktrace ⚠️

julia> solver = solve(prob, NewtonRaphson())
/AppleInternal/Library/BuildRoots/1c8f7852-1ca9-11f0-b28b-226177e5bb69/Library/Caches/com.apple.xbs/Sources/MetalPerformanceShaders/MPSMatrix/LinearAlgebra/MPSMatrixCombination.mm:293: failed assertion `[MPSMatrixCopy initWithDevice:subMatrixRows:subMatrixColumns] subMatrixColumns must be larger than zero'

[23292] signal 6: Abort trap: 6
in expression starting at REPL[5]:1
__pthread_kill at /usr/lib/system/libsystem_kernel.dylib (unknown line)
Allocations: 74535095 (Pool: 74532844; Big: 2251); GC: 38
[1]    23292 abort      julia

• Output of versioninfo()

Julia Version 1.11.4
Commit 8561cc3d68d (2025-03-10 11:36 UTC)
Build Info:
  Official https://julialang.org/ release
Platform Info:
  OS: macOS (arm64-apple-darwin24.0.0)
  CPU: 8 × Apple M3
  WORD_SIZE: 64
  LLVM: libLLVM-16.0.6 (ORCJIT, apple-m2)
Threads: 1 default, 0 interactive, 1 GC (on 4 virtual cores)