From 3c9cdccfcb4f390ccd13404b2301aa675f03f675 Mon Sep 17 00:00:00 2001
From: Manolis Papadakis <mpapadakis@nvidia.com>
Date: Tue, 20 Feb 2024 16:54:30 -0800
Subject: [PATCH] Basic np.load implementation, falling back to NumPy

Mostly there to ensure that the NumPy ndarray coming out of np.load
gets cast to a cuNumeric ndarray. As requested by ORNL.
---
 cunumeric/module.py | 52 ++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 51 insertions(+), 1 deletion(-)

diff --git a/cunumeric/module.py b/cunumeric/module.py
index c076200f8..0205c38f3 100644
--- a/cunumeric/module.py
+++ b/cunumeric/module.py
@@ -59,7 +59,8 @@
 from .utils import AxesPairLike, inner_modes, matmul_modes, tensordot_modes
 
 if TYPE_CHECKING:
-    from typing import Callable
+    from os import PathLike
+    from typing import BinaryIO, Callable
 
     import numpy.typing as npt
 
@@ -616,6 +617,55 @@ def copy(a: ndarray) -> ndarray:
     return result
 
 
+def load(
+    file: str | bytes | PathLike[Any] | BinaryIO,
+    *,
+    max_header_size: int = 10000,
+) -> ndarray:
+    """
+    Load an array from a ``.npy`` file.
+
+    Parameters
+    ----------
+    file : file-like object, string, or pathlib.Path
+        The file to read. File-like objects must support the
+        ``seek()`` and ``read()`` methods and must always
+        be opened in binary mode.
+    max_header_size : int, optional
+        Maximum allowed size of the header.  Large headers may not be safe
+        to load securely and thus require explicitly passing a larger value.
+        See :py:func:`ast.literal_eval()` for details.
+
+    Returns
+    -------
+    result : array
+        Data stored in the file.
+
+    Raises
+    ------
+    OSError
+        If the input file does not exist or cannot be read.
+
+    See Also
+    --------
+    numpy.load
+
+    Notes
+    -----
+    cuNumeric does not currently support ``.npz`` and pickled files.
+
+    Availability
+    --------
+    Single CPU
+    """
+    return array(
+        np.load(
+            file,
+            max_header_size=max_header_size,  # type: ignore [call-arg]
+        )
+    )
+
+
 # Numerical ranges