From 5b0f047226e482bdf94b0b09a91442ae3c0ccb6b Mon Sep 17 00:00:00 2001 From: CCYeh Date: Sat, 22 Nov 2025 20:09:44 +0100 Subject: [PATCH] Fix `mx.core.load` type annotation (#2819) --- python/src/ops.cpp | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/python/src/ops.cpp b/python/src/ops.cpp index c28cf6a51..25dd3cf51 100644 --- a/python/src/ops.cpp +++ b/python/src/ops.cpp @@ -4021,7 +4021,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]]"), + "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]]]"), R"pbdoc( Load array(s) from a binary file. @@ -4037,11 +4037,12 @@ void init_ops(nb::module_& m) { which support matadata. The metadata will be returned as an additional dictionary. Default: ``False``. Returns: - array or dict: + array, dict, or tuple: A single array if loading from a ``.npy`` file or a dict mapping names to arrays if loading from a ``.npz`` or - ``.safetensors`` file. If ``return_metadata`` is ``True`` an - additional dictionary of metadata will be returned. + ``.safetensors`` file. If ``return_metadata`` is ``True`` a + tuple ``(arrays, metadata)`` will be returned where the second + element is a dictionary containing the metadata. Warning: