[Augustas11]

• I understand it is strictly prohibited to use AI to write PR description

AI usage disclosure: drafted with AI assistance from a stream audit. Both call sites were checked against multiply and equal in mlx/ops.h, which take a stream. The test suite has not been run locally.

Follow-up to #4746. Two call sites still drop the stream the caller passed, and the stream checker misses them because they are operator overloads.

Complex l2_norm did abs(a, s) * abs(a, s). operator* takes no stream, so the product is built on the thread's default stream. Real inputs already go through square(a, s). The complex path now uses multiply(abs(a, s), abs(a, s), s).

Max and min Scatter::vjp did gathered_result == updates inside a multiply that already received stream(). operator== takes no stream. That comparison is now equal(gathered_result, updates, stream()).

This only changes complex l2_norm and the max/min scatter gradient.

The existing cross-thread test in tests/scheduler_tests.cpp now builds both graphs on a thread-unsafe stream and evaluates them from the other thread: a complex linalg::norm, and a scatter_max VJP.

Augustas from Malibu