Fix DeviceType annotations for device and stream arguments

[XXXXRT666]

Proposed changes

mx.cpu and mx.gpu are DeviceType values. Nanobind accepts them at runtime through the existing implicit conversion from DeviceType to Device, but the generated Python stubs only exposed Device. This caused type checkers to reject valid calls such as:

mx.exp(x, stream=mx.gpu)

This PR:

• Adds DeviceType as an explicit StreamOrDevice alternative.
• Updates stream argument annotations across core, random, linalg, fast, and distributed APIs.
• Includes ThreadLocalStream, which is already accepted by StreamOrDevice, in those annotations.
• Adds DeviceType-aware annotations to device and stream APIs that rely on nanobind's implicit conversion.

Checklist

• I have read the CONTRIBUTING document
• I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
• I have added tests that prove my fix is effective or that my feature works
• I have updated the necessary documentation (if needed)

[XXXXRT666]

Another question: Do you prefer using Union/Optional, or the Python 3.10-style A | B type annotation?

Because nanobind automatically generates the A | B syntax, while the existing codebase still uses the legacy Union/Optional style.

[zcbenz]

I don't really have a preference, but it feels more reasonable following what nanobind does.

[XXXXRT666]

I think the CI is failing because the non-Metal export cannot recognize DeviceType. Trying to fix it.