diff --git a/mlx/backend/common/metal_kernel.cpp b/mlx/backend/common/metal_kernel.cpp
index 4795f08778..43ed8c2047 100644
--- a/mlx/backend/common/metal_kernel.cpp
+++ b/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.");
diff --git a/mlx/utils.cpp b/mlx/utils.cpp
index ffe1e23173..4e29e8847b 100644
--- a/mlx/utils.cpp
+++ b/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);
   }
diff --git a/mlx/utils.h b/mlx/utils.h
index 9fa332fd2e..b5b516d89c 100644
--- a/mlx/utils.h
+++ b/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_);
 
diff --git a/python/mlx/_stub_patterns.txt b/python/mlx/_stub_patterns.txt
index c05f879d13..2a3fbce133 100644
--- a/python/mlx/_stub_patterns.txt
+++ b/python/mlx/_stub_patterns.txt
@@ -1,10 +1,5 @@
 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):
@@ -12,28 +7,27 @@ mlx.core.__prefix__:
     __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
diff --git a/python/src/array.cpp b/python/src/array.cpp
index 653b4a4c5b..6f65e14cde 100644
--- a/python/src/array.cpp
+++ b/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,
diff --git a/python/src/device.cpp b/python/src/device.cpp
index e70b69bd34..83d32d1cb6 100644
--- a/python/src/device.cpp
+++ b/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.
 
diff --git a/python/src/distributed.cpp b/python/src/distributed.cpp
index ed80001df5..0ad91c6e48 100644
--- a/python/src/distributed.cpp
+++ b/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.
diff --git a/python/src/export.cpp b/python/src/export.cpp
index 30d48bcae4..3596e4f846 100644
--- a/python/src/export.cpp
+++ b/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.
 
diff --git a/python/src/fast.cpp b/python/src/fast.cpp
index cd30b0bacd..e59357bc33 100644
--- a/python/src/fast.cpp
+++ b/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.
 
diff --git a/python/src/linalg.cpp b/python/src/linalg.cpp
index 0bf7b6f12c..a5e82ae853 100644
--- a/python/src/linalg.cpp
+++ b/python/src/linalg.cpp
@@ -55,7 +55,7 @@ void init_linalg(nb::module_& parent_module) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def norm(a: array, /, ord: Union[None, int, float, str] = None, axis: Union[None, int, list[int]] = None, keepdims: bool = False, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def norm(a: array, /, ord: None | int | float | str = None, axis: None | int | list[int] = None, keepdims: bool = False, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Matrix or vector norm.
 
@@ -177,7 +177,7 @@ void init_linalg(nb::module_& parent_module) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def qr(a: array, *, stream: Union[None, Stream, Device] = None) -> Tuple[array, array]"),
+          "def qr(a: array, *, stream: StreamOrDevice = None) -> tuple[array, array]"),
       R"pbdoc(
         The QR factorization of the input matrix.
 
@@ -220,7 +220,7 @@ void init_linalg(nb::module_& parent_module) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def svd(a: array, compute_uv: bool = True, *, stream: Union[None, Stream, Device] = None) -> Tuple[array, array, array]"),
+          "def svd(a: array, compute_uv: bool = True, *, stream: StreamOrDevice = None) -> tuple[array, array, array]"),
       R"pbdoc(
         The Singular Value Decomposition (SVD) of the input matrix.
 
@@ -246,8 +246,7 @@ void init_linalg(nb::module_& parent_module) {
       "a"_a,
       nb::kw_only(),
       "stream"_a = nb::none(),
-      nb::sig(
-          "def inv(a: array, *, stream: Union[None, Stream, Device] = None) -> array"),
+      nb::sig("def inv(a: array, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Compute the inverse of a square matrix.
 
@@ -271,7 +270,7 @@ void init_linalg(nb::module_& parent_module) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def tri_inv(a: array, upper: bool = False, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def tri_inv(a: array, upper: bool = False, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Compute the inverse of a triangular square matrix.
 
@@ -296,7 +295,7 @@ void init_linalg(nb::module_& parent_module) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def cholesky(a: array, upper: bool = False, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def cholesky(a: array, upper: bool = False, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Compute the Cholesky decomposition of a real symmetric positive semi-definite matrix.
 
@@ -326,7 +325,7 @@ void init_linalg(nb::module_& parent_module) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def cholesky_inv(a: array, upper: bool = False, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def cholesky_inv(a: array, upper: bool = False, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Compute the inverse of a real symmetric positive semi-definite matrix using it's Cholesky decomposition.
 
@@ -363,8 +362,7 @@ void init_linalg(nb::module_& parent_module) {
       "a"_a,
       nb::kw_only(),
       "stream"_a = nb::none(),
-      nb::sig(
-          "def pinv(a: array, *, stream: Union[None, Stream, Device] = None) -> array"),
+      nb::sig("def pinv(a: array, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Compute the (Moore-Penrose) pseudo-inverse of a matrix.
 
@@ -390,7 +388,7 @@ void init_linalg(nb::module_& parent_module) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def cross(a: array, b: array, axis: int = -1, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def cross(a: array, b: array, axis: int = -1, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Compute the cross product of two arrays along a specified axis.
 
@@ -449,7 +447,7 @@ void init_linalg(nb::module_& parent_module) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def eig(a: array, *, stream: Union[None, Stream, Device] = None) -> Tuple[array, array]"),
+          "def eig(a: array, *, stream: StreamOrDevice = None) -> tuple[array, array]"),
       R"pbdoc(
         Compute the eigenvalues and eigenvectors of a square matrix.
 
@@ -527,7 +525,7 @@ void init_linalg(nb::module_& parent_module) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def eigh(a: array, UPLO: str = 'L', *, stream: Union[None, Stream, Device] = None) -> Tuple[array, array]"),
+          "def eigh(a: array, UPLO: str = 'L', *, stream: StreamOrDevice = None) -> tuple[array, array]"),
       R"pbdoc(
         Compute the eigenvalues and eigenvectors of a complex Hermitian or
         real symmetric matrix.
@@ -573,7 +571,7 @@ void init_linalg(nb::module_& parent_module) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def lu(a: array, *, stream: Union[None, Stream, Device] = None) -> Tuple[array, array, array]"),
+          "def lu(a: array, *, stream: StreamOrDevice = None) -> tuple[array, array, array]"),
       R"pbdoc(
         Compute the LU factorization of the given matrix ``A``.
 
@@ -604,7 +602,7 @@ void init_linalg(nb::module_& parent_module) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def lu_factor(a: array, *, stream: Union[None, Stream, Device] = None) -> Tuple[array, array]"),
+          "def lu_factor(a: array, *, stream: StreamOrDevice = None) -> tuple[array, array]"),
       R"pbdoc(
         Computes a compact representation of the LU factorization.
 
@@ -624,7 +622,7 @@ void init_linalg(nb::module_& parent_module) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def solve(a: array, b: array, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def solve(a: array, b: array, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Compute the solution to a system of linear equations ``AX = B``.
 
@@ -646,7 +644,7 @@ void init_linalg(nb::module_& parent_module) {
       "upper"_a = false,
       "stream"_a = nb::none(),
       nb::sig(
-          "def solve_triangular(a: array, b: array, *, upper: bool = False, stream: Union[None, Stream, Device] = None) -> array"),
+          "def solve_triangular(a: array, b: array, *, upper: bool = False, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Computes the solution of a triangular system of linear equations ``AX = B``.
 
@@ -668,8 +666,7 @@ void init_linalg(nb::module_& parent_module) {
       "a"_a,
       nb::kw_only(),
       "stream"_a = nb::none(),
-      nb::sig(
-          "def det(a: array, *, stream: Union[None, Stream, Device] = None) -> array"),
+      nb::sig("def det(a: array, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Compute the determinant of a square matrix.
 
@@ -701,7 +698,7 @@ void init_linalg(nb::module_& parent_module) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def slogdet(a: array, *, stream: Union[None, Stream, Device] = None) -> Tuple[array, array]"),
+          "def slogdet(a: array, *, stream: StreamOrDevice = None) -> tuple[array, array]"),
       R"pbdoc(
         Compute the sign and natural log of the absolute value of the
         determinant of a square matrix.
diff --git a/python/src/ops.cpp b/python/src/ops.cpp
index 4e0373a238..8ddb3e4f96 100644
--- a/python/src/ops.cpp
+++ b/python/src/ops.cpp
@@ -63,7 +63,7 @@ void init_ops(nb::module_& m) {
       "stream"_a = nb::none(),
       nb::sig(
           "def reshape(a: array, /, shape: Sequence[int], *, stream: "
-          "Union[None, Stream, Device] = None) -> array"),
+          "StreamOrDevice = None) -> array"),
       R"pbdoc(
         Reshape an array while preserving the size.
 
@@ -91,7 +91,7 @@ void init_ops(nb::module_& m) {
       "stream"_a = nb::none(),
       nb::sig(
           "def flatten(a: array, /, start_axis: int = 0, end_axis: int = "
-          "-1, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "-1, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
       Flatten an array.
 
@@ -127,7 +127,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def unflatten(a: array, /, axis: int, shape: Sequence[int], *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def unflatten(a: array, /, axis: int, shape: Sequence[int], *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
       Unflatten an axis of an array to a shape.
 
@@ -164,8 +164,8 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def squeeze(a: array, /, axis: Union[None, int, Sequence[int]] = "
-          "None, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def squeeze(a: array, /, axis: None | int | Sequence[int] = "
+          "None, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Remove length one axes from an array.
 
@@ -193,8 +193,8 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def flip(a: array, /, axis: Union[None, int, Sequence[int]] = None, "
-          "*, stream: Union[None, Stream, Device] = None) -> array"),
+          "def flip(a: array, /, axis: None | int | Sequence[int] = None, "
+          "*, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Reverse the order of elements along the given axis.
 
@@ -216,8 +216,7 @@ void init_ops(nb::module_& m) {
       "axis"_a = 0,
       "stream"_a = nb::none(),
       nb::sig(
-          "def unstack(x: array, /, *, axis: int = 0, stream: Union[None, "
-          "Stream, Device] = None) -> list[array]"),
+          "def unstack(x: array, /, *, axis: int = 0, stream: StreamOrDevice = None) -> list[array]"),
       R"pbdoc(
         Split an array into a sequence of arrays along the given axis.
 
@@ -247,8 +246,8 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def expand_dims(a: array, /, axis: Union[int, Sequence[int]], "
-          "*, stream: Union[None, Stream, Device] = None) -> array"),
+          "def expand_dims(a: array, /, axis: int | Sequence[int], "
+          "*, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Add a size one dimension at the given axis.
 
@@ -268,7 +267,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def abs(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def abs(a: array, /, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Element-wise absolute value.
 
@@ -287,7 +286,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def sign(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def sign(a: array, /, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Element-wise sign.
 
@@ -304,7 +303,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def positive(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def positive(a: array, /, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Element-wise unary plus. Returns a copy of the input.
 
@@ -323,7 +322,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def negative(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def negative(a: array, /, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Element-wise negation.
 
@@ -346,7 +345,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def add(a: Union[scalar, array], b: Union[scalar, array], stream: Union[None, Stream, Device] = None) -> array"),
+          "def add(a: scalar | array, b: scalar | array, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Element-wise addition.
 
@@ -373,7 +372,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def subtract(a: Union[scalar, array], b: Union[scalar, array], stream: Union[None, Stream, Device] = None) -> array"),
+          "def subtract(a: scalar | array, b: scalar | array, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Element-wise subtraction.
 
@@ -400,7 +399,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def multiply(a: Union[scalar, array], b: Union[scalar, array], stream: Union[None, Stream, Device] = None) -> array"),
+          "def multiply(a: scalar | array, b: scalar | array, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Element-wise multiplication.
 
@@ -427,7 +426,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def divide(a: Union[scalar, array], b: Union[scalar, array], stream: Union[None, Stream, Device] = None) -> array"),
+          "def divide(a: scalar | array, b: scalar | array, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Element-wise division.
 
@@ -454,7 +453,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def divmod(a: Union[scalar, array], b: Union[scalar, array], stream: Union[None, Stream, Device] = None) -> array"),
+          "def divmod(a: scalar | array, b: scalar | array, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Element-wise quotient and remainder.
 
@@ -482,7 +481,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def floor_divide(a: Union[scalar, array], b: Union[scalar, array], stream: Union[None, Stream, Device] = None) -> array"),
+          "def floor_divide(a: scalar | array, b: scalar | array, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Element-wise integer division.
 
@@ -509,7 +508,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def remainder(a: Union[scalar, array], b: Union[scalar, array], stream: Union[None, Stream, Device] = None) -> array"),
+          "def remainder(a: scalar | array, b: scalar | array, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Element-wise remainder of division.
 
@@ -537,7 +536,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def equal(a: Union[scalar, array], b: Union[scalar, array], stream: Union[None, Stream, Device] = None) -> array"),
+          "def equal(a: scalar | array, b: scalar | array, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Element-wise equality.
 
@@ -564,7 +563,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def not_equal(a: Union[scalar, array], b: Union[scalar, array], stream: Union[None, Stream, Device] = None) -> array"),
+          "def not_equal(a: scalar | array, b: scalar | array, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Element-wise not equal.
 
@@ -591,7 +590,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def less(a: Union[scalar, array], b: Union[scalar, array], stream: Union[None, Stream, Device] = None) -> array"),
+          "def less(a: scalar | array, b: scalar | array, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Element-wise less than.
 
@@ -618,7 +617,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def less_equal(a: Union[scalar, array], b: Union[scalar, array], stream: Union[None, Stream, Device] = None) -> array"),
+          "def less_equal(a: scalar | array, b: scalar | array, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Element-wise less than or equal.
 
@@ -645,7 +644,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def greater(a: Union[scalar, array], b: Union[scalar, array], stream: Union[None, Stream, Device] = None) -> array"),
+          "def greater(a: scalar | array, b: scalar | array, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Element-wise greater than.
 
@@ -672,7 +671,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def greater_equal(a: Union[scalar, array], b: Union[scalar, array], stream: Union[None, Stream, Device] = None) -> array"),
+          "def greater_equal(a: scalar | array, b: scalar | array, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Element-wise greater or equal.
 
@@ -701,7 +700,7 @@ void init_ops(nb::module_& m) {
       "equal_nan"_a = false,
       "stream"_a = nb::none(),
       nb::sig(
-          "def array_equal(a: Union[scalar, array], b: Union[scalar, array], equal_nan: bool = False, stream: Union[None, Stream, Device] = None) -> array"),
+          "def array_equal(a: scalar | array, b: scalar | array, equal_nan: bool = False, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Array equality check.
 
@@ -726,7 +725,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def matmul(a: array, b: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def matmul(a: array, b: array, /, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Matrix multiplication.
 
@@ -757,7 +756,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def trunc(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def trunc(a: array, /, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Element-wise truncation towards zero.
 
@@ -776,7 +775,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def square(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def square(a: array, /, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Element-wise square.
 
@@ -795,7 +794,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def sqrt(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def sqrt(a: array, /, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Element-wise square root.
 
@@ -814,7 +813,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def rsqrt(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def rsqrt(a: array, /, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Element-wise reciprocal and square root.
 
@@ -833,7 +832,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def reciprocal(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def reciprocal(a: array, /, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Element-wise reciprocal.
 
@@ -852,7 +851,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def logical_not(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def logical_not(a: array, /, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Element-wise logical not.
 
@@ -872,7 +871,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def logical_and(a: array, b: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def logical_and(a: array, b: array, /, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Element-wise logical and.
 
@@ -894,7 +893,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def logical_or(a: array, b: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def logical_or(a: array, b: array, /, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Element-wise logical or.
 
@@ -915,7 +914,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def logical_xor(a: Union[scalar, array], b: Union[scalar, array], /, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def logical_xor(a: scalar | array, b: scalar | array, /, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Element-wise logical exclusive or.
 
@@ -939,7 +938,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def logaddexp(a: Union[scalar, array], b: Union[scalar, array], /, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def logaddexp(a: scalar | array, b: scalar | array, /, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Element-wise log-add-exp.
 
@@ -964,7 +963,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def exp(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def exp(a: array, /, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Element-wise exponential.
 
@@ -983,7 +982,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def expm1(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def expm1(a: array, /, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Element-wise exponential minus 1.
 
@@ -1004,7 +1003,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def erf(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def erf(a: array, /, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Element-wise error function.
 
@@ -1026,7 +1025,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def erfinv(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def erfinv(a: array, /, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Element-wise inverse of :func:`erf`.
 
@@ -1045,7 +1044,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def sin(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def sin(a: array, /, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Element-wise sine.
 
@@ -1064,7 +1063,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def cos(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def cos(a: array, /, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Element-wise cosine.
 
@@ -1083,7 +1082,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def tan(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def tan(a: array, /, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Element-wise tangent.
 
@@ -1102,7 +1101,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def arcsin(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def arcsin(a: array, /, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Element-wise inverse sine.
 
@@ -1121,7 +1120,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def arccos(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def arccos(a: array, /, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Element-wise inverse cosine.
 
@@ -1140,7 +1139,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def arctan(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def arctan(a: array, /, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Element-wise inverse tangent.
 
@@ -1158,7 +1157,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def arctan2(a: array, b: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def arctan2(a: array, b: array, /, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Element-wise inverse tangent of the ratio of two arrays.
 
@@ -1178,7 +1177,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def sinh(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def sinh(a: array, /, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Element-wise hyperbolic sine.
 
@@ -1197,7 +1196,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def cosh(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def cosh(a: array, /, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Element-wise hyperbolic cosine.
 
@@ -1216,7 +1215,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def tanh(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def tanh(a: array, /, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Element-wise hyperbolic tangent.
 
@@ -1235,7 +1234,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def arcsinh(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def arcsinh(a: array, /, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Element-wise inverse hyperbolic sine.
 
@@ -1254,7 +1253,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def arccosh(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def arccosh(a: array, /, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Element-wise inverse hyperbolic cosine.
 
@@ -1273,7 +1272,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def arctanh(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def arctanh(a: array, /, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Element-wise inverse hyperbolic tangent.
 
@@ -1292,7 +1291,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def degrees(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def degrees(a: array, /, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
       Convert angles from radians to degrees.
 
@@ -1311,7 +1310,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def radians(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def radians(a: array, /, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
       Convert angles from degrees to radians.
 
@@ -1330,7 +1329,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def log(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def log(a: array, /, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Element-wise natural logarithm.
 
@@ -1349,7 +1348,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def log2(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def log2(a: array, /, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Element-wise base-2 logarithm.
 
@@ -1368,7 +1367,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def log10(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def log10(a: array, /, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Element-wise base-10 logarithm.
 
@@ -1387,7 +1386,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def log1p(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def log1p(a: array, /, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Element-wise natural log of one plus the array.
 
@@ -1404,7 +1403,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def stop_gradient(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def stop_gradient(a: array, /, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Stop gradients from being computed.
 
@@ -1428,7 +1427,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def sigmoid(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def sigmoid(a: array, /, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Element-wise logistic sigmoid.
 
@@ -1456,7 +1455,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def power(a: Union[scalar, array], b: Union[scalar, array], /, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def power(a: scalar | array, b: scalar | array, /, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Element-wise power operation.
 
@@ -1503,7 +1502,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def arange(start : Union[int, float], stop : Union[None, int, float], step : Union[None, int, float], dtype: Optional[Dtype] = None, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def arange(start : int | float, stop : None | int | float, step : None | int | float, dtype: Dtype | None = None, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
       Generates ranges of numbers.
 
@@ -1548,7 +1547,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def arange(stop : Union[int, float], step : Union[None, int, float] = None, dtype: Optional[Dtype] = None, *, stream: Union[None, Stream, Device] = None) -> array"));
+          "def arange(stop : int | float, step : None | int | float = None, dtype: Dtype | None = None, *, stream: StreamOrDevice = None) -> array"));
   m.def(
       "bartlett",
       &mlx::core::bartlett,
@@ -1599,8 +1598,7 @@ void init_ops(nb::module_& m) {
       "M"_a,
       nb::kw_only(),
       "stream"_a = nb::none(),
-      nb::sig(
-          "def hamming(M: int, *, stream: Union[None, Stream, Device] = None) -> array"),
+      nb::sig("def hamming(M: int, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Return the Hamming window.
 
@@ -1624,7 +1622,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def blackman(M: int, *, stream: Union[None, Stream, Device] = None) -> array"), // <--- J'ai rajouté ça
+          "def blackman(M: int, *, stream: StreamOrDevice = None) -> array"), // <--- J'ai rajouté ça
       R"pbdoc(
         Return the Blackman window.
         
@@ -1661,7 +1659,7 @@ void init_ops(nb::module_& m) {
       "dtype"_a.none() = mx::float32,
       "stream"_a = nb::none(),
       nb::sig(
-          "def linspace(start: scalar, stop: scalar, num: Optional[int] = 50, dtype: Optional[Dtype] = float32, stream: Union[None, Stream, Device] = None) -> array"),
+          "def linspace(start: scalar, stop: scalar, num: int | None = 50, dtype: Dtype | None = float32, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Generate ``num`` evenly spaced numbers over interval ``[start, stop]``.
 
@@ -1683,14 +1681,14 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def kron(a: array, b: array, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def kron(a: array, b: array, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Compute the Kronecker product of two arrays ``a`` and ``b``.
 
         Args:
           a (array): The first input array.
           b (array): The second input array.
-          stream (Union[None, Stream, Device], optional): Optional stream or
+          stream (StreamOrDevice, optional): Optional stream or
             device for execution. Default: ``None``.
 
         Returns:
@@ -1727,7 +1725,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def take(a: array, /, indices: Union[int, array], axis: Optional[int] = None, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def take(a: array, /, indices: int | array, axis: int | None = None, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Take elements along an axis.
 
@@ -1764,7 +1762,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def take_along_axis(a: array, /, indices: array, axis: Optional[int] = None, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def take_along_axis(a: array, /, indices: array, axis: int | None = None, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Take values along an axis at the specified indices.
 
@@ -1803,7 +1801,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def put_along_axis(a: array, /, indices: array, values: array, axis: Optional[int] = None, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def put_along_axis(a: array, /, indices: array, values: array, axis: int | None = None, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Put values along an axis at the specified indices.
 
@@ -1835,7 +1833,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def full(shape: Union[int, Sequence[int]], vals: Union[scalar, array], dtype: Optional[Dtype] = None, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def full(shape: int | Sequence[int], vals: scalar | array, dtype: Dtype | None = None, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Construct an array with the given value.
 
@@ -1866,7 +1864,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def full_like(a: array, vals: Union[scalar, array], dtype: Optional[Dtype] = None, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def full_like(a: array, vals: scalar | array, dtype: Dtype | None = None, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         An array filled with ``vals`` with the same shape as the input.
 
@@ -1892,7 +1890,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def zeros(shape: Union[int, Sequence[int]], dtype: Optional[Dtype] = float32, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def zeros(shape: int | Sequence[int], dtype: Dtype | None = float32, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Construct an array of zeros.
 
@@ -1914,8 +1912,8 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "copy"_a = nb::none(),
       nb::sig(
-          "def asarray(a: Union[scalar, array, Sequence, DLPackCompatible], dtype: "
-          "Optional[Dtype] = None, *, copy: Optional[bool] = None) -> array"),
+          "def asarray(a: scalar | array | Sequence | DLPackCompatible, dtype: "
+          "Dtype | None = None, *, copy: bool | None = None) -> array"),
       R"pbdoc(
         Convert the input to an array.
 
@@ -1942,7 +1940,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "copy"_a = nb::none(),
       nb::sig(
-          "def from_dlpack(x: DLPackCompatible, /, *, copy: Optional[bool] = None) -> array"),
+          "def from_dlpack(x: DLPackCompatible, /, *, copy: bool | None = None) -> array"),
       R"pbdoc(
         Create an array from an object that supports DLPack.
 
@@ -1969,7 +1967,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def zeros_like(a: array, /, dtype: Optional[Dtype] = None, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def zeros_like(a: array, /, dtype: Dtype | None = None, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         An array of zeros like the input.
 
@@ -1994,7 +1992,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def ones(shape: Union[int, Sequence[int]], dtype: Optional[Dtype] = float32, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def ones(shape: int | Sequence[int], dtype: Dtype | None = float32, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Construct an array of ones.
 
@@ -2018,7 +2016,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def ones_like(a: array, /, dtype: Optional[Dtype] = None, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def ones_like(a: array, /, dtype: Dtype | None = None, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         An array of ones like the input.
 
@@ -2046,7 +2044,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def eye(n: int, m: Optional[int] = None, k: int = 0, dtype: Optional[Dtype] = float32, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def eye(n: int, m: int | None = None, k: int = 0, dtype: Dtype | None = float32, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Create an identity matrix or a general diagonal matrix.
 
@@ -2070,7 +2068,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def identity(n: int, dtype: Optional[Dtype] = float32, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def identity(n: int, dtype: Dtype | None = float32, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Create a square identity matrix.
 
@@ -2098,7 +2096,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def tri(n: int, m: int, k: int, dtype: Optional[Dtype] = None, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def tri(n: int, m: int, k: int, dtype: Dtype | None = None, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         An array with ones at and below the given diagonal and zeros elsewhere.
 
@@ -2120,7 +2118,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def tril(x: array, k: int, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def tril(x: array, k: int, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Zeros the array above the given diagonal.
 
@@ -2140,7 +2138,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def triu(x: array, k: int, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def triu(x: array, k: int, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Zeros the array below the given diagonal.
 
@@ -2163,7 +2161,7 @@ void init_ops(nb::module_& m) {
       "equal_nan"_a = false,
       "stream"_a = nb::none(),
       nb::sig(
-          "def allclose(a: array, b: array, /, rtol: float = 1e-05, atol: float = 1e-08, *, equal_nan: bool = False, stream: Union[None, Stream, Device] = None) -> array"),
+          "def allclose(a: array, b: array, /, rtol: float = 1e-05, atol: float = 1e-08, *, equal_nan: bool = False, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Approximate comparison of two arrays.
 
@@ -2200,7 +2198,7 @@ void init_ops(nb::module_& m) {
       "equal_nan"_a = false,
       "stream"_a = nb::none(),
       nb::sig(
-          "def isclose(a: array, b: array, /, rtol: float = 1e-05, atol: float = 1e-08, *, equal_nan: bool = False, stream: Union[None, Stream, Device] = None) -> array"),
+          "def isclose(a: array, b: array, /, rtol: float = 1e-05, atol: float = 1e-08, *, equal_nan: bool = False, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Returns a boolean array where two arrays are element-wise equal within a tolerance.
 
@@ -2241,7 +2239,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def all(a: array, /, axis: Union[None, int, Sequence[int]] = None, keepdims: bool = False, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def all(a: array, /, axis: None | int | Sequence[int] = None, keepdims: bool = False, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         An `and` reduction over the given axes.
 
@@ -2270,7 +2268,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def any(a: array, /, axis: Union[None, int, Sequence[int]] = None, keepdims: bool = False, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def any(a: array, /, axis: None | int | Sequence[int] = None, keepdims: bool = False, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         An `or` reduction over the given axes.
 
@@ -2298,7 +2296,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def minimum(a: Union[scalar, array], b: Union[scalar, array], /, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def minimum(a: scalar | array, b: scalar | array, /, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Element-wise minimum.
 
@@ -2325,7 +2323,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def maximum(a: Union[scalar, array], b: Union[scalar, array], /, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def maximum(a: scalar | array, b: scalar | array, /, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Element-wise maximum.
 
@@ -2348,7 +2346,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def floor(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def floor(a: array, /, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Element-wise floor.
 
@@ -2367,7 +2365,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def ceil(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def ceil(a: array, /, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Element-wise ceil.
 
@@ -2385,8 +2383,7 @@ void init_ops(nb::module_& m) {
       nb::arg(),
       nb::kw_only(),
       "stream"_a = nb::none(),
-      nb::sig(
-          "def isnan(a: array, stream: Union[None, Stream, Device] = None) -> array"),
+      nb::sig("def isnan(a: array, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Return a boolean array indicating which elements are NaN.
 
@@ -2404,8 +2401,7 @@ void init_ops(nb::module_& m) {
       nb::arg(),
       nb::kw_only(),
       "stream"_a = nb::none(),
-      nb::sig(
-          "def isinf(a: array, stream: Union[None, Stream, Device] = None) -> array"),
+      nb::sig("def isinf(a: array, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Return a boolean array indicating which elements are +/- inifnity.
 
@@ -2423,8 +2419,7 @@ void init_ops(nb::module_& m) {
       nb::arg(),
       nb::kw_only(),
       "stream"_a = nb::none(),
-      nb::sig(
-          "def isfinite(a: array, stream: Union[None, Stream, Device] = None) -> array"),
+      nb::sig("def isfinite(a: array, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Return a boolean array indicating which elements are finite.
 
@@ -2444,14 +2439,13 @@ void init_ops(nb::module_& m) {
       nb::arg(),
       nb::kw_only(),
       "stream"_a = nb::none(),
-      nb::sig(
-          "def isposinf(a: array, stream: Union[None, Stream, Device] = None) -> array"),
+      nb::sig("def isposinf(a: array, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Return a boolean array indicating which elements are positive infinity.
 
         Args:
             a (array): Input array.
-            stream (Union[None, Stream, Device]): Optional stream or device.
+            stream (StreamOrDevice): Optional stream or device.
 
         Returns:
             array: The boolean array indicating which elements are positive infinity.
@@ -2464,14 +2458,13 @@ void init_ops(nb::module_& m) {
       nb::arg(),
       nb::kw_only(),
       "stream"_a = nb::none(),
-      nb::sig(
-          "def isneginf(a: array, stream: Union[None, Stream, Device] = None) -> array"),
+      nb::sig("def isneginf(a: array, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Return a boolean array indicating which elements are negative infinity.
 
         Args:
             a (array): Input array.
-            stream (Union[None, Stream, Device]): Optional stream or device.
+            stream (StreamOrDevice): Optional stream or device.
 
         Returns:
             array: The boolean array indicating which elements are negative infinity.
@@ -2485,7 +2478,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def moveaxis(a: array, /, source: int, destination: int, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def moveaxis(a: array, /, source: int, destination: int, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Move an axis to a new position.
 
@@ -2506,7 +2499,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def swapaxes(a: array, /, axis1 : int, axis2: int, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def swapaxes(a: array, /, axis1 : int, axis2: int, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Swap two axes of an array.
 
@@ -2534,7 +2527,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def transpose(a: array, /, axes: Optional[Sequence[int]] = None, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def transpose(a: array, /, axes: Sequence[int] | None = None, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Transpose the dimensions of the array.
 
@@ -2562,7 +2555,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def permute_dims(a: array, /, axes: Optional[Sequence[int]] = None, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def permute_dims(a: array, /, axes: Sequence[int] | None = None, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         See :func:`transpose`.
       )pbdoc");
@@ -2580,7 +2573,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def sum(a: array, /, axis: Union[None, int, Sequence[int]] = None, keepdims: bool = False, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def sum(a: array, /, axis: None | int | Sequence[int] = None, keepdims: bool = False, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Sum reduce the array over the given axes.
 
@@ -2616,7 +2609,7 @@ void init_ops(nb::module_& m) {
       "keepdims"_a = false,
       "stream"_a = nb::none(),
       nb::sig(
-          "def count_nonzero(a: array, /, *, axis: Union[None, int, Sequence[int]] = None, keepdims: bool = False, stream: Union[None, Stream, Device] = None) -> array"),
+          "def count_nonzero(a: array, /, *, axis: None | int | Sequence[int] = None, keepdims: bool = False, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Count the number of non-zero elements along the given axis.
 
@@ -2644,7 +2637,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def prod(a: array, /, axis: Union[None, int, Sequence[int]] = None, keepdims: bool = False, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def prod(a: array, /, axis: None | int | Sequence[int] = None, keepdims: bool = False, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         An product reduction over the given axes.
 
@@ -2673,7 +2666,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def min(a: array, /, axis: Union[None, int, Sequence[int]] = None, keepdims: bool = False, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def min(a: array, /, axis: None | int | Sequence[int] = None, keepdims: bool = False, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         A `min` reduction over the given axes.
 
@@ -2702,7 +2695,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def max(a: array, /, axis: Union[None, int, Sequence[int]] = None, keepdims: bool = False, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def max(a: array, /, axis: None | int | Sequence[int] = None, keepdims: bool = False, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         A `max` reduction over the given axes.
 
@@ -2738,7 +2731,7 @@ void init_ops(nb::module_& m) {
       "inclusive"_a = true,
       "stream"_a = nb::none(),
       nb::sig(
-          "def logcumsumexp(a: array, /, axis: Optional[int] = None, *, reverse: bool = False, inclusive: bool = True, stream: Union[None, Stream, Device] = None) -> array"),
+          "def logcumsumexp(a: array, /, axis: int | None = None, *, reverse: bool = False, inclusive: bool = True, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Return the cumulative logsumexp of the elements along the given axis.
 
@@ -2768,7 +2761,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def logsumexp(a: array, /, axis: Union[None, int, Sequence[int]] = None, keepdims: bool = False, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def logsumexp(a: array, /, axis: None | int | Sequence[int] = None, keepdims: bool = False, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         A `log-sum-exp` reduction over the given axes.
 
@@ -2803,7 +2796,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def mean(a: array, /, axis: Union[None, int, Sequence[int]] = None, keepdims: bool = False, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def mean(a: array, /, axis: None | int | Sequence[int] = None, keepdims: bool = False, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Compute the mean(s) over the given axes.
 
@@ -2832,7 +2825,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def median(a: array, /, axis: Union[None, int, Sequence[int]] = None, keepdims: bool = False, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def median(a: array, /, axis: None | int | Sequence[int] = None, keepdims: bool = False, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Compute the median(s) over the given axes.
 
@@ -2863,7 +2856,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def var(a: array, /, axis: Union[None, int, Sequence[int]] = None, keepdims: bool = False, ddof: int = 0, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def var(a: array, /, axis: None | int | Sequence[int] = None, keepdims: bool = False, ddof: int = 0, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Compute the variance(s) over the given axes.
 
@@ -2896,7 +2889,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def std(a: array, /, axis: Union[None, int, Sequence[int]] = None, keepdims: bool = False, ddof: int = 0, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def std(a: array, /, axis: None | int | Sequence[int] = None, keepdims: bool = False, ddof: int = 0, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Compute the standard deviation(s) over the given axes.
 
@@ -2932,7 +2925,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def split(a: array, /, indices_or_sections: Union[int, Sequence[int]], axis: int = 0, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def split(a: array, /, indices_or_sections: int | Sequence[int], axis: int = 0, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Split an array along a given axis.
 
@@ -2976,7 +2969,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def argmin(a: array, /, axis: Union[None, int] = None, keepdims: bool = False, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def argmin(a: array, /, axis: None | int = None, keepdims: bool = False, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Indices of the minimum values along the axis.
 
@@ -3008,7 +3001,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def argmax(a: array, /, axis: Union[None, int] = None, keepdims: bool = False, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def argmax(a: array, /, axis: None | int = None, keepdims: bool = False, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Indices of the maximum values along the axis.
 
@@ -3036,7 +3029,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def sort(a: array, /, axis: Union[None, int] = -1, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def sort(a: array, /, axis: None | int = -1, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Returns a sorted copy of the array.
 
@@ -3066,7 +3059,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def argsort(a: array, /, axis: Union[None, int] = -1, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def argsort(a: array, /, axis: None | int = -1, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Returns the indices that sort the array.
 
@@ -3100,7 +3093,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def partition(a: array, /, kth: int, axis: Union[None, int] = -1, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def partition(a: array, /, kth: int, axis: None | int = -1, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Returns a partitioned copy of the array such that the smaller ``kth``
         elements are first.
@@ -3138,7 +3131,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def argpartition(a: array, /, kth: int, axis: Union[None, int] = -1, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def argpartition(a: array, /, kth: int, axis: None | int = -1, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Returns the indices that partition the array.
 
@@ -3177,7 +3170,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def topk(a: array, /, k: int, axis: Union[None, int] = -1, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def topk(a: array, /, k: int, axis: None | int = -1, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Returns the ``k`` largest elements from the input along a given axis.
 
@@ -3203,7 +3196,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def broadcast_to(a: Union[scalar, array], /, shape: Sequence[int], *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def broadcast_to(a: scalar | array, /, shape: Sequence[int], *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Broadcast an array to the given shape.
 
@@ -3225,7 +3218,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def broadcast_arrays(*arrays: array, stream: Union[None, Stream, Device] = None) -> Tuple[array, ...]"),
+          "def broadcast_arrays(*arrays: array, stream: StreamOrDevice = None) -> tuple[array, ...]"),
       R"pbdoc(
         Broadcast arrays against one another.
 
@@ -3251,7 +3244,7 @@ void init_ops(nb::module_& m) {
       "precise"_a = false,
       "stream"_a = nb::none(),
       nb::sig(
-          "def softmax(a: array, /, axis: Union[None, int, Sequence[int]] = None, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def softmax(a: array, /, axis: None | int | Sequence[int] = None, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Perform the softmax along the given axis.
 
@@ -3286,7 +3279,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def concatenate(arrays: list[array], axis: Optional[int] = 0, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def concatenate(arrays: list[array], axis: int | None = 0, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Concatenate the arrays along the given axis.
 
@@ -3314,7 +3307,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def concat(arrays: list[array], axis: Optional[int] = 0, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def concat(arrays: list[array], axis: int | None = 0, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         See :func:`concatenate`.
       )pbdoc");
@@ -3334,7 +3327,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def stack(arrays: list[array], axis: Optional[int] = 0, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def stack(arrays: list[array], axis: int | None = 0, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Stacks the arrays along a new axis.
 
@@ -3362,7 +3355,7 @@ void init_ops(nb::module_& m) {
       "indexing"_a = "xy",
       "stream"_a = nb::none(),
       nb::sig(
-          "def meshgrid(*arrays: array, sparse: Optional[bool] = False, indexing: Optional[str] = 'xy', stream: Union[None, Stream, Device] = None) -> array"),
+          "def meshgrid(*arrays: array, sparse: bool | None = False, indexing: str | None = 'xy', stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Generate multidimensional coordinate grids from 1-D coordinate arrays
 
@@ -3395,7 +3388,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def repeat(array: array, repeats: int, axis: Optional[int] = None, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def repeat(array: array, repeats: int, axis: int | None = None, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Repeat an array along a specified axis.
 
@@ -3432,7 +3425,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def clip(a: array, /, a_min: Union[scalar, array, None], a_max: Union[scalar, array, None], *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def clip(a: array, /, a_min: scalar | array | None, a_max: scalar | array | None, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Clip the values of the array between the given minimum and maximum.
 
@@ -3482,7 +3475,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def pad(a: array, pad_width: Union[int, tuple[int], tuple[int, int], list[tuple[int, int]]], mode: Literal['constant', 'edge'] = 'constant', constant_values: Union[scalar, array] = 0, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def pad(a: array, pad_width: int | tuple[int] | tuple[int, int] | list[tuple[int, int]], mode: Literal['constant', 'edge'] = 'constant', constant_values: scalar | array = 0, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Pad an array with a constant value
 
@@ -3529,7 +3522,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def as_strided(a: array, /, shape: Optional[Sequence[int]] = None, strides: Optional[Sequence[int]] = None, offset: int = 0, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def as_strided(a: array, /, shape: Sequence[int] | None = None, strides: Sequence[int] | None = None, offset: int = 0, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Create a view into the array with the given shape and strides.
 
@@ -3566,7 +3559,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def astype(a: array, dtype: Dtype, /, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def astype(a: array, dtype: Dtype, /, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Cast the array to a specified type.
 
@@ -3598,7 +3591,7 @@ void init_ops(nb::module_& m) {
       "dtype"_a = nb::none(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def cumsum(a: array, /, axis: Optional[int] = None, *, reverse: bool = False, inclusive: bool = True, dtype: Optional[Dtype] = None, stream: Union[None, Stream, Device] = None) -> array"),
+          "def cumsum(a: array, /, axis: int | None = None, *, reverse: bool = False, inclusive: bool = True, dtype: Dtype | None = None, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Return the cumulative sum of the elements along the given axis.
 
@@ -3636,7 +3629,7 @@ void init_ops(nb::module_& m) {
       "dtype"_a = nb::none(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def cumprod(a: array, /, axis: Optional[int] = None, *, reverse: bool = False, inclusive: bool = True, dtype: Optional[Dtype] = None, stream: Union[None, Stream, Device] = None) -> array"),
+          "def cumprod(a: array, /, axis: int | None = None, *, reverse: bool = False, inclusive: bool = True, dtype: Dtype | None = None, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Return the cumulative product of the elements along the given axis.
 
@@ -3673,7 +3666,7 @@ void init_ops(nb::module_& m) {
       "inclusive"_a = true,
       "stream"_a = nb::none(),
       nb::sig(
-          "def cummax(a: array, /, axis: Optional[int] = None, *, reverse: bool = False, inclusive: bool = True, stream: Union[None, Stream, Device] = None) -> array"),
+          "def cummax(a: array, /, axis: int | None = None, *, reverse: bool = False, inclusive: bool = True, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Return the cumulative maximum of the elements along the given axis.
 
@@ -3709,7 +3702,7 @@ void init_ops(nb::module_& m) {
       "inclusive"_a = true,
       "stream"_a = nb::none(),
       nb::sig(
-          "def cummin(a: array, /, axis: Optional[int] = None, *, reverse: bool = False, inclusive: bool = True, stream: Union[None, Stream, Device] = None) -> array"),
+          "def cummin(a: array, /, axis: int | None = None, *, reverse: bool = False, inclusive: bool = True, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Return the cumulative minimum of the elements along the given axis.
 
@@ -3734,7 +3727,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def diff(a: array, /, n: int = 1, axis: int = -1, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def diff(a: array, /, n: int = 1, axis: int = -1, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         The n-th discrete difference along the given axis.
 
@@ -3755,8 +3748,7 @@ void init_ops(nb::module_& m) {
       nb::arg(),
       nb::kw_only(),
       "stream"_a = nb::none(),
-      nb::sig(
-          "def conj(a: array, *, stream: Union[None, Stream, Device] = None) -> array"),
+      nb::sig("def conj(a: array, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Return the elementwise complex conjugate of the input.
         Alias for `mx.conjugate`.
@@ -3776,7 +3768,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def conjugate(a: array, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def conjugate(a: array, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Return the elementwise complex conjugate of the input.
         Alias for `mx.conj`.
@@ -3850,7 +3842,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          R"(def convolve(a: array, v: array, /, mode: str = "full", *, stream: Union[None, Stream, Device] = None) -> array)"),
+          R"(def convolve(a: array, v: array, /, mode: str = "full", *, stream: StreamOrDevice = None) -> array)"),
       R"pbdoc(
         The discrete convolution of 1D arrays.
 
@@ -3877,7 +3869,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def conv1d(input: array, weight: array, /, stride: int = 1, padding: int = 0, dilation: int = 1, groups: int = 1, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def conv1d(input: array, weight: array, /, stride: int = 1, padding: int = 0, dilation: int = 1, groups: int = 1, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         1D convolution over an input with several channels
 
@@ -3935,7 +3927,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def conv2d(input: array, weight: array, /, stride: Union[int, tuple[int, int]] = 1, padding: Union[int, tuple[int, int]] = 0, dilation: Union[int, tuple[int, int]] = 1, groups: int = 1, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def conv2d(input: array, weight: array, /, stride: int | tuple[int, int] = 1, padding: int | tuple[int, int] = 0, dilation: int | tuple[int, int] = 1, groups: int = 1, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         2D convolution over an input with several channels
 
@@ -4005,7 +3997,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def conv3d(input: array, weight: array, /, stride: Union[int, tuple[int, int, int]] = 1, padding: Union[int, tuple[int, int, int]] = 0, dilation: Union[int, tuple[int, int, int]] = 1, groups: int = 1, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def conv3d(input: array, weight: array, /, stride: int | tuple[int, int, int] = 1, padding: int | tuple[int, int, int] = 0, dilation: int | tuple[int, int, int] = 1, groups: int = 1, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         3D convolution over an input with several channels
 
@@ -4041,7 +4033,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def conv_transpose1d(input: array, weight: array, /, stride: int = 1, padding: int = 0, dilation: int = 1, output_padding: int = 0, groups: int = 1, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def conv_transpose1d(input: array, weight: array, /, stride: int = 1, padding: int = 0, dilation: int = 1, output_padding: int = 0, groups: int = 1, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         1D transposed convolution over an input with several channels
 
@@ -4116,7 +4108,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def conv_transpose2d(input: array, weight: array, /, stride: Union[int, Tuple[int, int]] = 1, padding: Union[int, Tuple[int, int]] = 0, dilation: Union[int, Tuple[int, int]] = 1, output_padding: Union[int, Tuple[int, int]] = 0, groups: int = 1, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def conv_transpose2d(input: array, weight: array, /, stride: int | tuple[int, int] = 1, padding: int | tuple[int, int] = 0, dilation: int | tuple[int, int] = 1, output_padding: int | tuple[int, int] = 0, groups: int = 1, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         2D transposed convolution over an input with several channels
 
@@ -4202,7 +4194,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def conv_transpose3d(input: array, weight: array, /, stride: Union[int, Tuple[int, int, int]] = 1, padding: Union[int, Tuple[int, int, int]] = 0, dilation: Union[int, Tuple[int, int, int]] = 1, output_padding: Union[int, Tuple[int, int, int]] = 0, groups: int = 1, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def conv_transpose3d(input: array, weight: array, /, stride: int | tuple[int, int, int] = 1, padding: int | tuple[int, int, int] = 0, dilation: int | tuple[int, int, int] = 1, output_padding: int | tuple[int, int, int] = 0, groups: int = 1, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         3D transposed convolution over an input with several channels
 
@@ -4304,7 +4296,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def conv_general(input: array, weight: array, /, stride: Union[int, Sequence[int]] = 1, padding: Union[int, Sequence[int], tuple[Sequence[int], Sequence[int]]] = 0, kernel_dilation: Union[int, Sequence[int]] = 1, input_dilation: Union[int, Sequence[int]] = 1, groups: int = 1, flip: bool = False, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def conv_general(input: array, weight: array, /, stride: int | Sequence[int] = 1, padding: int | Sequence[int] | tuple[Sequence[int], Sequence[int]] = 0, kernel_dilation: int | Sequence[int] = 1, input_dilation: int | Sequence[int] = 1, groups: int = 1, flip: bool = False, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         General convolution over an input with several channels
 
@@ -4337,8 +4329,7 @@ void init_ops(nb::module_& m) {
       &mlx_save_helper,
       "file"_a,
       "arr"_a,
-      nb::sig(
-          "def save(file: Union[file, str, pathlib.Path], arr: array) -> None"),
+      nb::sig("def save(file: file | str | pathlib.Path, arr: array) -> None"),
       R"pbdoc(
         Save the array to a binary file in ``.npy`` format.
 
@@ -4354,8 +4345,7 @@ void init_ops(nb::module_& m) {
       "file"_a,
       "args"_a,
       "kwargs"_a,
-      nb::sig(
-          "def savez(file: Union[file, str, pathlib.Path], *args, **kwargs)"),
+      nb::sig("def savez(file: file | str | pathlib.Path, *args, **kwargs)"),
       R"pbdoc(
         Save several arrays to a binary file in uncompressed ``.npz``
         format.
@@ -4389,7 +4379,7 @@ void init_ops(nb::module_& m) {
       "args"_a,
       "kwargs"_a,
       nb::sig(
-          "def savez_compressed(file: Union[file, str, pathlib.Path], *args, **kwargs)"),
+          "def savez_compressed(file: file | str | pathlib.Path, *args, **kwargs)"),
       R"pbdoc(
         Save several arrays to a binary file in compressed ``.npz`` format.
 
@@ -4408,7 +4398,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def load(file: Union[file, str, pathlib.Path], /, format: Optional[str] = None, return_metadata: bool = False, *, stream: Union[None, Stream, Device] = None) -> Union[array, dict[str, array], Tuple[dict[str, array], dict[str, Any]]]"),
+          "def load(file: file | str | pathlib.Path, /, format: str | None = None, return_metadata: bool = False, *, stream: StreamOrDevice = None) -> array | dict[str, array] | tuple[dict[str, array], dict[str, Any]]"),
       R"pbdoc(
         Load array(s) from a binary file.
 
@@ -4443,7 +4433,7 @@ void init_ops(nb::module_& m) {
       "arrays"_a,
       "metadata"_a = nb::none(),
       nb::sig(
-          "def save_safetensors(file: Union[file, str, pathlib.Path], arrays: dict[str, array], metadata: Optional[dict[str, str]] = None)"),
+          "def save_safetensors(file: file | str | pathlib.Path, arrays: dict[str, array], metadata: dict[str, str] | None = None)"),
       R"pbdoc(
         Save array(s) to a binary file in ``.safetensors`` format.
 
@@ -4465,7 +4455,7 @@ void init_ops(nb::module_& m) {
       "arrays"_a,
       "metadata"_a = nb::none(),
       nb::sig(
-          "def save_gguf(file: Union[file, str, pathlib.Path], arrays: dict[str, array], metadata: dict[str, Union[array, str, list[str]]])"),
+          "def save_gguf(file: file | str | pathlib.Path, arrays: dict[str, array], metadata: dict[str, array | str | list[str]])"),
       R"pbdoc(
         Save array(s) to a binary file in ``.gguf`` format.
 
@@ -4496,7 +4486,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def where(condition: Union[scalar, array], x: Union[scalar, array], y: Union[scalar, array], /, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def where(condition: scalar | array, x: scalar | array, y: scalar | array, /, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Select from ``x`` or ``y`` according to ``condition``.
 
@@ -4528,7 +4518,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def nan_to_num(a: Union[scalar, array], nan: float = 0, posinf: Optional[float] = None, neginf: Optional[float] = None, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def nan_to_num(a: scalar | array, nan: float = 0, posinf: float | None = None, neginf: float | None = None, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Replace NaN and Inf values with finite numbers.
 
@@ -4555,7 +4545,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def round(a: array, /, decimals: int = 0, stream: Union[None, Stream, Device] = None) -> array"),
+          "def round(a: array, /, decimals: int = 0, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Round to the given number of decimals.
 
@@ -4588,7 +4578,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def quantized_matmul(x: array, w: array, /, scales: array, biases: Optional[array] = None, transpose: bool = True, group_size: Optional[int] = None, bits: Optional[int] = None, mode: str = 'affine', *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def quantized_matmul(x: array, w: array, /, scales: array, biases: array | None = None, transpose: bool = True, group_size: int | None = None, bits: int | None = None, mode: str = 'affine', *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Perform the matrix multiplication with the quantized matrix ``w``. The
         quantization uses one floating point scale and bias per ``group_size`` of
@@ -4626,7 +4616,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def quantize(w: array, /, group_size: Optional[int] = None, bits: Optional[int] = None, mode: str = 'affine', *, global_scale: Optional[array] = None, stream: Union[None, Stream, Device] = None) -> tuple[array, array, array]"),
+          "def quantize(w: array, /, group_size: int | None = None, bits: int | None = None, mode: str = 'affine', *, global_scale: array | None = None, stream: StreamOrDevice = None) -> tuple[array, array, array]"),
       R"pbdoc(
         Quantize the array ``w``.
 
@@ -4727,7 +4717,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def dequantize(w: array, /, scales: array, biases: Optional[array] = None, group_size: Optional[int] = None, bits: Optional[int] = None, mode: str = 'affine', global_scale: Optional[array] = None, dtype: Optional[Dtype] = None, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def dequantize(w: array, /, scales: array, biases: array | None = None, group_size: int | None = None, bits: int | None = None, mode: str = 'affine', global_scale: array | None = None, dtype: Dtype | None = None, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Dequantize the matrix ``w`` using quantization parameters.
 
@@ -4782,7 +4772,7 @@ void init_ops(nb::module_& m) {
       "sorted_indices"_a = false,
       "stream"_a = nb::none(),
       nb::sig(
-          "def gather_qmm(x: array, w: array, /, scales: array, biases: Optional[array] = None, lhs_indices: Optional[array] = None, rhs_indices: Optional[array] = None, transpose: bool = True, group_size: Optional[int] = None, bits: Optional[int] = None, mode: str = 'affine', *, sorted_indices: bool = False, stream: Union[None, Stream, Device] = None) -> array"),
+          "def gather_qmm(x: array, w: array, /, scales: array, biases: array | None = None, lhs_indices: array | None = None, rhs_indices: array | None = None, transpose: bool = True, group_size: int | None = None, bits: int | None = None, mode: str = 'affine', *, sorted_indices: bool = False, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Perform quantized matrix multiplication with matrix-level gather.
 
@@ -4836,7 +4826,7 @@ void init_ops(nb::module_& m) {
       "sorted_indices"_a = false,
       "stream"_a = nb::none(),
       nb::sig(
-          "def gather_qqmm(x: array, w: array, /, scales: Optional[array] = None, lhs_indices: Optional[array] = None, rhs_indices: Optional[array] = None, group_size: Optional[int] = None, bits: Optional[int] = None, mode: str = 'nvfp4', global_scale_x: Optional[array] = None, global_scale_w: Optional[array] = None, *, sorted_indices: bool = False, stream: Union[None, Stream, Device] = None) -> array"),
+          "def gather_qqmm(x: array, w: array, /, scales: array | None = None, lhs_indices: array | None = None, rhs_indices: array | None = None, group_size: int | None = None, bits: int | None = None, mode: str = 'nvfp4', global_scale_x: array | None = None, global_scale_w: array | None = None, *, sorted_indices: bool = False, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Fused :func:`qqmm` with matrix-level gather.
 
@@ -4880,7 +4870,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def segmented_mm(a: array, b: array, /, segments: array, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def segmented_mm(a: array, b: array, /, segments: array, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Perform a matrix multiplication but segment the inner dimension and
         save the result for each segment separately.
@@ -4916,7 +4906,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def tensordot(a: array, b: array, /, axes: Union[int, list[Sequence[int]]] = 2, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def tensordot(a: array, b: array, /, axes: int | list[Sequence[int]] = 2, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Compute the tensor dot product along the specified axes.
 
@@ -4940,7 +4930,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def inner(a: array, b: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def inner(a: array, b: array, /, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
       Ordinary inner product of vectors for 1-D arrays, in higher dimensions a sum product over the last axes.
 
@@ -4960,7 +4950,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def vecdot(a: array, b: array, /, *, axis: int = -1, stream: Union[None, Stream, Device] = None) -> array"),
+          "def vecdot(a: array, b: array, /, *, axis: int = -1, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
       Compute the vector dot product of two arrays along an axis.
 
@@ -4980,7 +4970,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def outer(a: array, b: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def outer(a: array, b: array, /, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
       Compute the outer product of two 1-D arrays, if the array's passed are not 1-D a flatten op will be run beforehand.
 
@@ -5007,7 +4997,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def tile(a: array, reps: Union[int, Sequence[int]], /, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def tile(a: array, reps: int | Sequence[int], /, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
       Construct an array by repeating ``a`` the number of times given by ``reps``.
 
@@ -5029,7 +5019,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def addmm(c: array, a: array, b: array, /, alpha: float = 1.0, beta: float = 1.0,  *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def addmm(c: array, a: array, b: array, /, alpha: float = 1.0, beta: float = 1.0,  *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Matrix multiplication with addition and optional scaling.
 
@@ -5059,7 +5049,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def block_masked_mm(a: array, b: array, /, block_size: int = 64, mask_out: Optional[array] = None, mask_lhs: Optional[array] = None, mask_rhs: Optional[array] = None, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def block_masked_mm(a: array, b: array, /, block_size: int = 64, mask_out: array | None = None, mask_lhs: array | None = None, mask_rhs: array | None = None, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Matrix multiplication with block masking.
 
@@ -5098,7 +5088,7 @@ void init_ops(nb::module_& m) {
       "sorted_indices"_a = false,
       "stream"_a = nb::none(),
       nb::sig(
-          "def gather_mm(a: array, b: array, /, lhs_indices: array, rhs_indices: array, *, sorted_indices: bool = False, stream: Union[None, Stream, Device] = None) -> array"),
+          "def gather_mm(a: array, b: array, /, lhs_indices: array, rhs_indices: array, *, sorted_indices: bool = False, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Matrix multiplication with matrix-level gather.
 
@@ -5140,7 +5130,7 @@ void init_ops(nb::module_& m) {
       "axis2"_a = 1,
       "stream"_a = nb::none(),
       nb::sig(
-          "def diagonal(a: array, offset: int = 0, axis1: int = 0, axis2: int = 1, stream: Union[None, Stream, Device] = None) -> array"),
+          "def diagonal(a: array, offset: int = 0, axis1: int = 0, axis2: int = 1, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Return specified diagonals.
 
@@ -5172,7 +5162,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def diag(a: array, /, k: int = 0, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def diag(a: array, /, k: int = 0, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Extract a diagonal or construct a diagonal matrix.
         If ``a`` is 1-D then a diagonal matrix is constructed with ``a`` on the
@@ -5208,7 +5198,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def trace(a: array, /, offset: int = 0, axis1: int = 0, axis2: int = 1, dtype: Optional[Dtype] = None, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def trace(a: array, /, offset: int = 0, axis1: int = 0, axis2: int = 1, dtype: Dtype | None = None, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Return the sum along a specified diagonal in the given array.
 
@@ -5238,13 +5228,13 @@ void init_ops(nb::module_& m) {
       "arys"_a,
       "stream"_a = nb::none(),
       nb::sig(
-          "def atleast_1d(*arys: array, stream: Union[None, Stream, Device] = None) -> Union[array, list[array]]"),
+          "def atleast_1d(*arys: array, stream: StreamOrDevice = None) -> array | list[array]"),
       R"pbdoc(
         Convert all arrays to have at least one dimension.
 
         Args:
             *arys: Input arrays.
-            stream (Union[None, Stream, Device], optional): The stream to execute the operation on.
+            stream (StreamOrDevice, optional): The stream to execute the operation on.
 
         Returns:
             array or list(array): An array or list of arrays with at least one dimension.
@@ -5261,13 +5251,13 @@ void init_ops(nb::module_& m) {
       "arys"_a,
       "stream"_a = nb::none(),
       nb::sig(
-          "def atleast_2d(*arys: array, stream: Union[None, Stream, Device] = None) -> Union[array, list[array]]"),
+          "def atleast_2d(*arys: array, stream: StreamOrDevice = None) -> array | list[array]"),
       R"pbdoc(
         Convert all arrays to have at least two dimensions.
 
         Args:
             *arys: Input arrays.
-            stream (Union[None, Stream, Device], optional): The stream to execute the operation on.
+            stream (StreamOrDevice, optional): The stream to execute the operation on.
 
         Returns:
             array or list(array): An array or list of arrays with at least two dimensions.
@@ -5284,13 +5274,13 @@ void init_ops(nb::module_& m) {
       "arys"_a,
       "stream"_a = nb::none(),
       nb::sig(
-          "def atleast_3d(*arys: array, stream: Union[None, Stream, Device] = None) -> Union[array, list[array]]"),
+          "def atleast_3d(*arys: array, stream: StreamOrDevice = None) -> array | list[array]"),
       R"pbdoc(
         Convert all arrays to have at least three dimensions.
 
         Args:
             *arys: Input arrays.
-            stream (Union[None, Stream, Device], optional): The stream to execute the operation on.
+            stream (StreamOrDevice, optional): The stream to execute the operation on.
 
         Returns:
             array or list(array): An array or list of arrays with at least three dimensions.
@@ -5320,7 +5310,7 @@ void init_ops(nb::module_& m) {
       ""_a,
       ""_a,
       nb::sig(
-          "def issubdtype(arg1: Union[Dtype, DtypeCategory], arg2: Union[Dtype, DtypeCategory]) -> bool"),
+          "def issubdtype(arg1: Dtype | DtypeCategory, arg2: Dtype | DtypeCategory) -> bool"),
       R"pbdoc(
         Check if a :obj:`Dtype` or :obj:`DtypeCategory` is a subtype
         of another.
@@ -5392,8 +5382,7 @@ void init_ops(nb::module_& m) {
         }
         return t;
       },
-      nb::sig(
-          "def result_type(*arrays_and_dtypes: Union[array, Dtype]) -> Dtype"),
+      nb::sig("def result_type(*arrays_and_dtypes: array | Dtype) -> Dtype"),
       R"pbdoc(
         The type that results from applying type promotion to the inputs.
 
@@ -5420,7 +5409,7 @@ void init_ops(nb::module_& m) {
       },
       "from_"_a,
       "to"_a,
-      nb::sig("def can_cast(from_: Union[array, Dtype], to: Dtype) -> bool"),
+      nb::sig("def can_cast(from_: array | Dtype, to: Dtype) -> bool"),
       R"pbdoc(
         Determine if one data type can be cast to another according to type
         promotion rules.
@@ -5481,7 +5470,7 @@ void init_ops(nb::module_& m) {
       "dtype"_a,
       "kind"_a,
       nb::sig(
-          "def isdtype(dtype: Dtype, kind: Union[Dtype, str, tuple[Union[Dtype, str], ...]]) -> bool"),
+          "def isdtype(dtype: Dtype, kind: Dtype | str | tuple[Dtype | str, ...]) -> bool"),
       R"pbdoc(
         Test whether a dtype belongs to one or more data type kinds.
 
@@ -5509,7 +5498,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def bitwise_and(a: Union[scalar, array], b: Union[scalar, array], stream: Union[None, Stream, Device] = None) -> array"),
+          "def bitwise_and(a: scalar | array, b: scalar | array, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Element-wise bitwise and.
 
@@ -5536,7 +5525,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def bitwise_or(a: Union[scalar, array], b: Union[scalar, array], stream: Union[None, Stream, Device] = None) -> array"),
+          "def bitwise_or(a: scalar | array, b: scalar | array, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Element-wise bitwise or.
 
@@ -5563,7 +5552,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def bitwise_xor(a: Union[scalar, array], b: Union[scalar, array], stream: Union[None, Stream, Device] = None) -> array"),
+          "def bitwise_xor(a: scalar | array, b: scalar | array, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Element-wise bitwise xor.
 
@@ -5591,7 +5580,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def left_shift(a: Union[scalar, array], b: Union[scalar, array], stream: Union[None, Stream, Device] = None) -> array"),
+          "def left_shift(a: scalar | array, b: scalar | array, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Element-wise left shift.
 
@@ -5619,7 +5608,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def right_shift(a: Union[scalar, array], b: Union[scalar, array], stream: Union[None, Stream, Device] = None) -> array"),
+          "def right_shift(a: scalar | array, b: scalar | array, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Element-wise right shift.
 
@@ -5644,7 +5633,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def bitwise_invert(a: Union[scalar, array], stream: Union[None, Stream, Device] = None) -> array"),
+          "def bitwise_invert(a: scalar | array, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Element-wise bitwise inverse.
 
@@ -5666,7 +5655,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def view(a: Union[scalar, array], dtype: Dtype, stream: Union[None, Stream, Device] = None) -> array"),
+          "def view(a: scalar | array, dtype: Dtype, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         View the array as a different type.
 
@@ -5692,7 +5681,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def hadamard_transform(a: array, scale: Optional[float] = None, stream: Union[None, Stream, Device] = None) -> array"),
+          "def hadamard_transform(a: array, scale: float | None = None, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Perform the Walsh-Hadamard transform along the final axis.
 
@@ -5756,7 +5745,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def einsum(subscripts: str, *operands, stream: Union[None, Stream, Device] = None) -> array"),
+          "def einsum(subscripts: str, *operands, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
 
       Perform the Einstein summation convention on the operands.
@@ -5791,7 +5780,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def roll(a: array, shift: Union[int, Tuple[int]], axis: Union[None, int, Tuple[int]] = None, /, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def roll(a: array, shift: int | tuple[int], axis: None | int | tuple[int] = None, /, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Roll array elements along a given axis.
 
@@ -5819,7 +5808,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def real(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def real(a: array, /, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Returns the real part of a complex array.
 
@@ -5838,7 +5827,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def imag(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def imag(a: array, /, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Returns the imaginary part of a complex array.
 
@@ -5865,7 +5854,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def slice(a: array, start_indices: array, axes: Sequence[int], slice_size: Sequence[int], *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def slice(a: array, start_indices: array, axes: Sequence[int], slice_size: Sequence[int], *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Extract a sub-array from the input array.
 
@@ -5904,7 +5893,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def slice_update(a: array, update: array, start_indices: array, axes: Sequence[int], *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def slice_update(a: array, update: array, start_indices: array, axes: Sequence[int], *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Update a sub-array of the input array.
 
@@ -5933,7 +5922,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def contiguous(a: array, /, allow_col_major: bool = False, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def contiguous(a: array, /, allow_col_major: bool = False, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
       Force an array to be row contiguous. Copy if necessary.
 
@@ -5962,7 +5951,7 @@ void init_ops(nb::module_& m) {
 
         return nb::tuple(nb::cast(result));
       },
-      nb::sig("def broadcast_shapes(*shapes: Sequence[int]) -> Tuple[int]"),
+      nb::sig("def broadcast_shapes(*shapes: Sequence[int]) -> tuple[int]"),
       R"pbdoc(
         Broadcast shapes.
 
@@ -6013,7 +6002,7 @@ void init_ops(nb::module_& m) {
       nb::arg(),
       nb::arg(),
       nb::sig(
-          "def depends(inputs: Union[array, Sequence[array]], dependencies: Union[array, Sequence[array]])"),
+          "def depends(inputs: array | Sequence[array], dependencies: array | Sequence[array])"),
       R"pbdoc(
         Insert dependencies between arrays in the graph. The outputs are
         identical to ``inputs`` but with dependencies on ``dependencies``.
@@ -6040,7 +6029,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def qqmm(x: array, w: array, scales: Optional[array] = None, group_size: Optional[int] = None, bits: Optional[int] = None, mode: str = 'nvfp4', global_scale_x: Optional[array] = None, global_scale_w: Optional[array] = None, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def qqmm(x: array, w: array, scales: array | None = None, group_size: int | None = None, bits: int | None = None, mode: str = 'nvfp4', global_scale_x: array | None = None, global_scale_w: array | None = None, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
       Perform a matrix multiplication using a possibly quantized weight matrix
       ``w`` and a non-quantized input ``x``. The input ``x`` is quantized on the
@@ -6090,7 +6079,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def from_fp8(x: array, dtype: Dtype = bfloat16, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def from_fp8(x: array, dtype: Dtype = bfloat16, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
       Convert the array from fp8 (e4m3) to another floating-point type.
 
@@ -6108,7 +6097,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def to_fp8(x: array, *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def to_fp8(x: array, *, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
       Convert the array to fp8 (e4m3) from another floating-point type.
 
diff --git a/python/src/random.cpp b/python/src/random.cpp
index ceebb52e7c..66e407a1e0 100644
--- a/python/src/random.cpp
+++ b/python/src/random.cpp
@@ -149,7 +149,7 @@ void init_random(nb::module_& parent_module) {
       "num"_a = 2,
       "stream"_a = nb::none(),
       nb::sig(
-          "def split(key: array, num: int = 2, stream: Union[None, Stream, Device] = None) -> array"),
+          "def split(key: array, num: int = 2, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Split a PRNG key into sub keys.
 
@@ -184,7 +184,7 @@ void init_random(nb::module_& parent_module) {
       "key"_a = nb::none(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def uniform(low: Union[scalar, array] = 0, high: Union[scalar, array] = 1, shape: Sequence[int] = [], dtype: Optional[Dtype] = float32, key: Optional[array] = None, stream: Union[None, Stream, Device] = None) -> array"),
+          "def uniform(low: scalar | array = 0, high: scalar | array = 1, shape: Sequence[int] = [], dtype: Dtype | None = float32, key: array | None = None, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Generate uniformly distributed random numbers.
 
@@ -227,7 +227,7 @@ void init_random(nb::module_& parent_module) {
       "key"_a = nb::none(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def normal(shape: Sequence[int] = [], dtype: Optional[Dtype] = float32, loc: Union[scalar, array, None] = None, scale: Union[scalar, array, None] = None, key: Optional[array] = None, stream: Union[None, Stream, Device] = None) -> array"),
+          "def normal(shape: Sequence[int] = [], dtype: Dtype | None = float32, loc: scalar | array | None = None, scale: scalar | array | None = None, key: array | None = None, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Generate normally distributed random numbers.
 
@@ -267,7 +267,7 @@ void init_random(nb::module_& parent_module) {
       "key"_a = nb::none(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def multivariate_normal(mean: array, cov: array, shape: Sequence[int] = [], dtype: Optional[Dtype] = float32, key: Optional[array] = None, stream: Union[None, Stream, Device] = None) -> array"),
+          "def multivariate_normal(mean: array, cov: array, shape: Sequence[int] = [], dtype: Dtype | None = float32, key: array | None = None, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Generate jointly-normal random samples given a mean and covariance.
 
@@ -314,7 +314,7 @@ void init_random(nb::module_& parent_module) {
       "key"_a = nb::none(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def randint(low: Union[scalar, array], high: Union[scalar, array], shape: Sequence[int] = [], dtype: Optional[Dtype] = int32, key: Optional[array] = None, stream: Union[None, Stream, Device] = None) -> array"),
+          "def randint(low: scalar | array, high: scalar | array, shape: Sequence[int] = [], dtype: Dtype | None = int32, key: array | None = None, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Generate random integers from the given interval.
 
@@ -351,7 +351,7 @@ void init_random(nb::module_& parent_module) {
       "key"_a = nb::none(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def bernoulli(p: Union[scalar, array] = 0.5, shape: Optional[Sequence[int]] = None, key: Optional[array] = None, stream: Union[None, Stream, Device] = None) -> array"),
+          "def bernoulli(p: scalar | array = 0.5, shape: Sequence[int] | None = None, key: array | None = None, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Generate Bernoulli random values.
 
@@ -395,7 +395,7 @@ void init_random(nb::module_& parent_module) {
       "key"_a = nb::none(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def truncated_normal(lower: Union[scalar, array], upper: Union[scalar, array], shape: Optional[Sequence[int]] = None, dtype: Optional[Dtype] = float32, key: Optional[array] = None, stream: Union[None, Stream, Device] = None) -> array"),
+          "def truncated_normal(lower: scalar | array, upper: scalar | array, shape: Sequence[int] | None = None, dtype: Dtype | None = float32, key: array | None = None, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Generate values from a truncated normal distribution.
 
@@ -429,7 +429,7 @@ void init_random(nb::module_& parent_module) {
       "key"_a = nb::none(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def gumbel(shape: Sequence[int] = [], dtype: Optional[Dtype] = float32, key: Optional[array] = None, stream: Union[None, Stream, Device] = None) -> array"),
+          "def gumbel(shape: Sequence[int] = [], dtype: Dtype | None = float32, key: array | None = None, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Sample from the standard Gumbel distribution.
 
@@ -475,7 +475,7 @@ void init_random(nb::module_& parent_module) {
       "key"_a = nb::none(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def categorical(logits: array, axis: int = -1, shape: Optional[Sequence[int]] = None, num_samples: Optional[int] = None, key: Optional[array] = None, stream: Union[None, Stream, Device] = None) -> array"),
+          "def categorical(logits: array, axis: int = -1, shape: Sequence[int] | None = None, num_samples: int | None = None, key: array | None = None, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Sample from a categorical distribution.
 
@@ -518,7 +518,7 @@ void init_random(nb::module_& parent_module) {
       "key"_a = nb::none(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def laplace(shape: Sequence[int] = [], dtype: Optional[Dtype] = float32, loc: float = 0.0, scale: float = 1.0, key: Optional[array] = None, stream: Union[None, Stream, Device] = None) -> array"),
+          "def laplace(shape: Sequence[int] = [], dtype: Dtype | None = float32, loc: float = 0.0, scale: float = 1.0, key: array | None = None, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Sample numbers from a Laplace distribution.
 
@@ -551,7 +551,7 @@ void init_random(nb::module_& parent_module) {
       "key"_a = nb::none(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def permutation(x: Union[int, array], axis: int = 0, key: Optional[array] = None, stream: Union[None, Stream, Device] = None) -> array"),
+          "def permutation(x: int | array, axis: int = 0, key: array | None = None, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         Generate a random permutation or permute the entries of an array.
 
diff --git a/python/src/stream.cpp b/python/src/stream.cpp
index 76fe427696..004301a45c 100644
--- a/python/src/stream.cpp
+++ b/python/src/stream.cpp
@@ -88,6 +88,7 @@ void init_stream(nb::module_& m) {
       "default_stream",
       &mx::default_stream,
       "device"_a,
+      nb::sig("def default_stream(device: Device | DeviceType) -> Stream"),
       R"pbdoc(Get the device's default stream.)pbdoc");
   m.def(
       "set_default_stream",
@@ -106,6 +107,7 @@ void init_stream(nb::module_& m) {
       "new_stream",
       &mx::new_stream,
       "device"_a,
+      nb::sig("def new_stream(device: Device | DeviceType) -> Stream"),
       R"pbdoc(
         Make a new stream on the given device.
 
@@ -116,6 +118,8 @@ void init_stream(nb::module_& m) {
       "new_thread_unsafe_stream",
       &mx::new_thread_unsafe_stream,
       "device"_a,
+      nb::sig(
+          "def new_thread_unsafe_stream(device: Device | DeviceType) -> Stream"),
       R"pbdoc(
         Make a new stream that can be used in any thread.
 
@@ -128,6 +132,8 @@ void init_stream(nb::module_& m) {
       "new_thread_local_stream",
       &mx::new_thread_local_stream,
       "device"_a,
+      nb::sig(
+          "def new_thread_local_stream(device: Device | DeviceType) -> ThreadLocalStream"),
       R"pbdoc(Make a new stream that will be unique per thread.)pbdoc");
   m.def(
       "clear_streams",
diff --git a/python/src/transforms.cpp b/python/src/transforms.cpp
index 6e3fda8882..1d7aa8b9b1 100644
--- a/python/src/transforms.cpp
+++ b/python/src/transforms.cpp
@@ -1352,7 +1352,7 @@ void init_transforms(nb::module_& m) {
       "argnums"_a = nb::none(),
       "argnames"_a = std::vector<std::string>{},
       nb::sig(
-          "def value_and_grad(fun: Callable[P, R], argnums: Optional[Union[int, Sequence[int]]] = None, argnames: Union[str, Sequence[str]] = []) -> Callable[P, Tuple[R, Any]]"),
+          "def value_and_grad(fun: Callable[P, R], argnums: int | Sequence[int] | None = None, argnames: str | Sequence[str] = []) -> Callable[P, tuple[R, Any]]"),
       R"pbdoc(
         Returns a function which computes the value and gradient of ``fun``.
 
@@ -1421,7 +1421,7 @@ void init_transforms(nb::module_& m) {
       "argnums"_a = nb::none(),
       "argnames"_a = std::vector<std::string>{},
       nb::sig(
-          "def grad(fun: Callable[P, R], argnums: Optional[Union[int, Sequence[int]]] = None, argnames: Union[str, Sequence[str]] = []) -> Callable[P, Any]"),
+          "def grad(fun: Callable[P, R], argnums: int | Sequence[int] | None = None, argnames: str | Sequence[str] = []) -> Callable[P, Any]"),
       R"pbdoc(
         Returns a function which computes the gradient of ``fun``.
 
@@ -1491,7 +1491,7 @@ void init_transforms(nb::module_& m) {
       "outputs"_a = nb::none(),
       "shapeless"_a = false,
       nb::sig(
-          "def compile(fun: Callable[P, R], inputs: Optional[object] = None, outputs: Optional[object] = None, shapeless: bool = False) -> Callable[P, R]"),
+          "def compile(fun: Callable[P, R], inputs: object | None = None, outputs: object | None = None, shapeless: bool = False) -> Callable[P, R]"),
       R"pbdoc(
         Returns a compiled function which produces the same output as ``fun``.