Fix DeviceType annotations for device and stream arguments

mlx/backend/common/metal_kernel.cpp

@@ -37,9 +37,11 @@ Stream resolve_metal_kernel_stream(StreamOrDevice s) {
   // recorded in the graph on a placeholder GPU stream. The importing process
   // remaps it to one of its own streams.
   auto* device = std::get_if<Device>(&s);
+  auto* device_type = std::get_if<Device::DeviceType>(&s);
   auto* stream = std::get_if<Stream>(&s);
   auto* tl_stream = std::get_if<ThreadLocalStream>(&s);
   if ((device && *device != Device::gpu) ||
+      (device_type && *device_type != Device::gpu) ||
       (stream && stream->device != Device::gpu) ||
       (tl_stream && tl_stream->device != Device::gpu)) {
     throw std::invalid_argument("[metal_kernel] Only supports the GPU.");

mlx/utils.cpp

@@ -18,6 +18,8 @@ Stream to_stream(StreamOrDevice s) {
     return default_stream(default_device());
   } else if (std::holds_alternative<Device>(s)) {
     return default_stream(std::get<Device>(s));
+  } else if (std::holds_alternative<Device::DeviceType>(s)) {
+    return default_stream(std::get<Device::DeviceType>(s));
   } else if (std::holds_alternative<ThreadLocalStream>(s)) {
     return stream_from_thread_local_stream(std::get<ThreadLocalStream>(s));
   } else {
@@ -30,6 +32,10 @@ Stream to_stream(StreamOrDevice s, Device default_) {
     return default_stream(default_);
   } else if (std::holds_alternative<Device>(s)) {
     return default_stream(std::get<Device>(s));
+  } else if (std::holds_alternative<Device::DeviceType>(s)) {
+    return default_stream(std::get<Device::DeviceType>(s));
+  } else if (std::holds_alternative<ThreadLocalStream>(s)) {
+    return stream_from_thread_local_stream(std::get<ThreadLocalStream>(s));
   } else {
     return std::get<Stream>(s);
   }

mlx/utils.h

@@ -17,8 +17,12 @@
 
 namespace mlx::core {
 
-using StreamOrDevice =
-    std::variant<std::monostate, Stream, ThreadLocalStream, Device>;
+using StreamOrDevice = std::variant<
+    std::monostate,
+    Stream,
+    ThreadLocalStream,
+    Device,
+    Device::DeviceType>;
 MLX_API Stream to_stream(StreamOrDevice s);
 MLX_API Stream to_stream(StreamOrDevice s, Device default_);
 

python/mlx/_stub_patterns.txt

@@ -1,39 +1,33 @@
 mlx.core.__prefix__:
-  from typing import Any, Callable, Dict, List, Optional, Protocol, Sequence, Tuple, Union, ParamSpec, TypeVar
-  import sys
-  if sys.version_info >= (3, 10):
-    from typing import TypeAlias
-  else:
-    from typing_extensions import TypeAlias
+  from typing import Any, ParamSpec, Protocol, TypeAlias, TypeVar
   P = ParamSpec("P")
   R = TypeVar("R")
   class DLPackCompatible(Protocol):
     __dlpack__: Callable[..., Any]
     __dlpack_device__: Callable[..., Any]
 
 mlx.core.__suffix__:
-  from typing import Union
-  scalar: TypeAlias = Union[int, float, bool]
-  list_or_scalar: TypeAlias = Union[scalar, list["list_or_scalar"]]
+  scalar: TypeAlias = int | float | bool
+  list_or_scalar: TypeAlias = scalar | list["list_or_scalar"]
+  StreamOrDevice: TypeAlias = Stream | ThreadLocalStream | Device | DeviceType | None
   bool_: Dtype = ...
 
 mlx.core.distributed.__prefix__:
-  from mlx.core import array, Dtype, Device, Stream, scalar
+  from mlx.core import array, Dtype, StreamOrDevice, scalar
   from mlx.core.distributed import Group
-  from typing import Sequence, Optional, Union
+  from collections.abc import Sequence
 
 mlx.core.fast.__prefix__:
-  from mlx.core import array, Dtype, Device, Stream, scalar
-  from typing import Sequence, Optional, Union
+  from mlx.core import array, Dtype, StreamOrDevice, scalar
 
 mlx.core.linalg.__prefix__:
-  from mlx.core import array, Dtype, Device, Stream, scalar
-  from typing import Sequence, Optional, Tuple, Union
+  from mlx.core import array, Dtype, StreamOrDevice, scalar
+  from collections.abc import Sequence
 
 mlx.core.metal.__prefix__:
   from mlx.core import array, Dtype, Device, Stream, scalar
-  from typing import Sequence, Optional, Union
+  from collections.abc import Sequence
 
 mlx.core.random.__prefix__:
-  from mlx.core import array, Dtype, Device, Stream, scalar, float32, int32
-  from typing import Sequence, Optional, Union
+  from mlx.core import array, Dtype, StreamOrDevice, scalar, float32, int32
+  from collections.abc import Sequence

python/src/array.cpp

@@ -312,7 +312,7 @@ void init_array(nb::module_& m) {
           "val"_a,
           "dtype"_a = nb::none(),
           nb::sig(
-              "def __init__(self: array, val: Union[scalar, list, tuple, DLPackCompatible, array], dtype: Optional[Dtype] = None)"))
+              "def __init__(self: array, val: scalar | list | tuple | DLPackCompatible | array, dtype: Dtype | None = None)"))
       .def_prop_ro(
           "size",
           &mx::array::size,

python/src/device.cpp

@@ -61,11 +61,13 @@ void init_device(nb::module_& m) {
       "set_default_device",
       &mx::set_default_device,
       "device"_a,
+      nb::sig("def set_default_device(device: Device | DeviceType) -> None"),
       R"pbdoc(Set the default device.)pbdoc");
   m.def(
       "is_available",
       &mx::is_available,
       "device"_a,
+      nb::sig("def is_available(device: Device | DeviceType) -> bool"),
       R"pbdoc(Check if a back-end is available for the given device.)pbdoc");
   m.def(
       "device_count",
@@ -86,6 +88,8 @@ void init_device(nb::module_& m) {
         return mx::device_info(d.value_or(mx::default_device()));
       },
       "d"_a = nb::none(),
+      nb::sig(
+          "def device_info(d: None | Device | DeviceType = None) -> dict[str, str | int]"),
       R"pbdoc(
       Get information about a device.
 

python/src/distributed.cpp

@@ -125,7 +125,7 @@ void init_distributed(nb::module_& parent_module) {
       nb::kw_only(),
       "all_gather_factory"_a = nb::none(),
       nb::sig(
-          "def init(strict: bool = False, backend: str = 'any', *, all_gather_factory: Optional[Callable[[int, int], Callable[[bytes, int], bytes]]] = None) -> Group"),
+          "def init(strict: bool = False, backend: str = 'any', *, all_gather_factory: Callable[[int, int], Callable[[bytes, int], bytes]] | None = None) -> Group"),
       R"pbdoc(
         Initialize the communication backend and create the global communication group.
 
@@ -170,7 +170,7 @@ void init_distributed(nb::module_& parent_module) {
       "group"_a = nb::none(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def all_sum(x: array, *, group: Optional[Group] = None, stream: Union[None, Stream, Device] = None) -> array"),
+          "def all_sum(x: array, *, group: Group | None = None, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         All reduce sum.
 
@@ -199,7 +199,7 @@ void init_distributed(nb::module_& parent_module) {
       "group"_a = nb::none(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def all_max(x: array, *, group: Optional[Group] = None, stream: Union[None, Stream, Device] = None) -> array"),
+          "def all_max(x: array, *, group: Group | None = None, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         All reduce max.
 
@@ -228,7 +228,7 @@ void init_distributed(nb::module_& parent_module) {
       "group"_a = nb::none(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def all_min(x: array, *, group: Optional[Group] = None, stream: Union[None, Stream, Device] = None) -> array"),
+          "def all_min(x: array, *, group: Group | None = None, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
       All reduce min.
 
@@ -257,7 +257,7 @@ void init_distributed(nb::module_& parent_module) {
       "group"_a = nb::none(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def all_gather(x: array, *, group: Optional[Group] = None, stream: Union[None, Stream, Device] = None) -> array"),
+          "def all_gather(x: array, *, group: Group | None = None, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Gather arrays from all processes.
 
@@ -290,7 +290,7 @@ void init_distributed(nb::module_& parent_module) {
       "group"_a = nb::none(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def send(x: array, dst: int, *, group: Optional[Group] = None, stream: Union[None, Stream, Device] = None) -> array"),
+          "def send(x: array, dst: int, *, group: Group | None = None, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Send an array from the current process to the process that has rank
         ``dst`` in the group.
@@ -318,7 +318,7 @@ void init_distributed(nb::module_& parent_module) {
       "group"_a = nb::none(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def recv(shape: Sequence[int], dtype: Dtype, src: int, *, group: Optional[Group] = None, stream: Union[None, Stream, Device] = None) -> array"),
+          "def recv(shape: Sequence[int], dtype: Dtype, src: int, *, group: Group | None = None, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Recv an array with shape ``shape`` and dtype ``dtype`` from process
         with rank ``src``.
@@ -351,7 +351,7 @@ void init_distributed(nb::module_& parent_module) {
       "group"_a = nb::none(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def recv_like(x: array, src: int, *, group: Optional[Group] = None, stream: Union[None, Stream, Device] = None) -> array"),
+          "def recv_like(x: array, src: int, *, group: Group | None = None, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Recv an array with shape and type like ``x`` from process with rank
         ``src``.
@@ -384,7 +384,7 @@ void init_distributed(nb::module_& parent_module) {
       "group"_a = nb::none(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def sum_scatter(x: array, *, group: Optional[Group] = None, stream: Union[None, Stream, Device] = None) -> array"),
+          "def sum_scatter(x: array, *, group: Group | None = None, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
       Sum ``x`` across all processes in the group and shard the result along the first axis across ranks.
       ``x.shape[0]`` must be divisible by the group size.

python/src/export.cpp

@@ -173,7 +173,7 @@ void init_export(nb::module_& m) {
       "metadata"_a = nb::none(),
       "kwargs"_a,
       nb::sig(
-          "def export_function(file_or_callback: Union[str, Callable], fun: Callable, *args, shapeless: bool = False, metadata: Optional[str] = None, **kwargs) -> None"),
+          "def export_function(file_or_callback: str | Callable, fun: Callable, *args, shapeless: bool = False, metadata: str | None = None, **kwargs) -> None"),
       R"pbdoc(
         Export an MLX function.
 
@@ -236,7 +236,7 @@ void init_export(nb::module_& m) {
       "file"_a,
       "return_metadata"_a = false,
       nb::sig(
-          "def import_function(file: str, return_metadata: bool = False) -> Union[Callable, tuple[Callable, str]]"),
+          "def import_function(file: str, return_metadata: bool = False) -> Callable | tuple[Callable, str]"),
       R"pbdoc(
         Import a function from a file.
 

python/src/fast.cpp

@@ -127,7 +127,7 @@ void init_fast(nb::module_& parent_module) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def rms_norm(x: array, weight: Optional[array], eps: float, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def rms_norm(x: array, weight: array | None, eps: float, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Root Mean Square normalization (RMS norm).
 
@@ -154,7 +154,7 @@ void init_fast(nb::module_& parent_module) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def layer_norm(x: array, weight: Optional[array], bias: Optional[array], eps: float, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def layer_norm(x: array, weight: array | None, bias: array | None, eps: float, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Layer normalization.
 
@@ -197,7 +197,7 @@ void init_fast(nb::module_& parent_module) {
       "freqs"_a = nb::none(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def rope(a: array, dims: int, *, traditional: bool, base: Optional[float], scale: float, offset: Union[int, array], freqs: Optional[array] = None, stream: Union[None, Stream, Device] = None) -> array"),
+          "def rope(a: array, dims: int, *, traditional: bool, base: float | None, scale: float, offset: int | array, freqs: array | None = None, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Apply rotary positional encoding to the input.
 
@@ -271,7 +271,7 @@ void init_fast(nb::module_& parent_module) {
       "sinks"_a = nb::none(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def scaled_dot_product_attention(q: array, k: array, v: array, *, scale: float,  mask: Union[None, str, array] = None, sinks: Optional[array] = None, stream: Union[None, Stream, Device] = None) -> array"),
+          "def scaled_dot_product_attention(q: array, k: array, v: array, *, scale: float,  mask: None | str | array = None, sinks: array | None = None, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         A fast implementation of multi-head attention: ``O = softmax(Q @ K.T, dim=-1) @ V``.
 
@@ -365,7 +365,7 @@ void init_fast(nb::module_& parent_module) {
             "verbose"_a = false,
             "stream"_a = nb::none(),
             nb::sig(
-                "def __call__(self, *, inputs: List[Union[scalar, array]], output_shapes: List[Sequence[int]], output_dtypes: List[Dtype], grid: tuple[int, int, int], threadgroup: tuple[int, int, int], template: Optional[List[Tuple[str, Union[bool, int, Dtype]]]] = None, init_value: Optional[float] = None, verbose: bool = false, stream: Union[None, Stream, Device] = None)"),
+                "def __call__(self, *, inputs: list[scalar | array], output_shapes: list[Sequence[int]], output_dtypes: list[Dtype], grid: tuple[int, int, int], threadgroup: tuple[int, int, int], template: list[tuple[str, bool | int | Dtype]] | None = None, init_value: float | None = None, verbose: bool = false, stream: StreamOrDevice = None)"),
             R"pbdoc(
             Run the kernel.
 
@@ -489,7 +489,7 @@ void init_fast(nb::module_& parent_module) {
             "verbose"_a = false,
             "stream"_a = nb::none(),
             nb::sig(
-                "def __call__(self, *, inputs: List[Union[scalar, array]], output_shapes: List[Sequence[int]], output_dtypes: List[Dtype], grid: tuple[int, int, int], threadgroup: tuple[int, int, int], template: Optional[List[Tuple[str, Union[bool, int, Dtype]]]] = None, init_value: Optional[float] = None, verbose: bool = false, stream: Union[None, Stream, Device] = None)"),
+                "def __call__(self, *, inputs: list[scalar | array], output_shapes: list[Sequence[int]], output_dtypes: list[Dtype], grid: tuple[int, int, int], threadgroup: tuple[int, int, int], template: list[tuple[str, bool | int | Dtype]] | None = None, init_value: float | None = None, verbose: bool = false, stream: StreamOrDevice = None)"),
             R"pbdoc(
             Run the kernel.