Repository: rsokl/MyGrad (language: python)

Issue:
Save and load func
Finished the save and load functions, thanks for the help by the way!


Discussion / hints:

Fix patch:
diff --git a/src/mygrad/__init__.py b/src/mygrad/__init__.py
index 6a323e63..fec278c2 100644
--- a/src/mygrad/__init__.py
+++ b/src/mygrad/__init__.py
@@ -32,6 +32,7 @@
 from mygrad.tensor_manip.tiling.funcs import *
 from mygrad.tensor_manip.transpose_like.funcs import *
 from mygrad.ufuncs._ufunc_creators import ufunc
+from ._io import load, save
 
 from . import random
 from ._version import get_versions
diff --git a/src/mygrad/_io.py b/src/mygrad/_io.py
new file mode 100644
index 00000000..ad80ac7f
--- /dev/null
+++ b/src/mygrad/_io.py
@@ -0,0 +1,125 @@
+from pathlib import Path
+from typing import BinaryIO, Union
+
+import numpy as np
+
+import mygrad.tensor_base as tb
+
+_FileLike = Union[str, Path, BinaryIO]
+
+
+def save(file: _FileLike, tensor: tb.Tensor) -> None:
+    """Saves a tensor and its gradient information.
+
+    This docstring was adapted from that of numpy.save()
+
+    Parameters
+    ----------
+    file : str | Path | BinaryIO
+        The file or file-path that where the tensor data and its gradient
+        will be saved. Note that the file will be saved as a .npz file.
+
+    tensor : Tensor
+        The tensor to be saved. If it has an associated gradient, that will
+        be saved as well.
+
+    Notes
+    -----
+    This function uses ``numpy.savez(file, data=tensor.data, grad=tensor.grad)``
+    to save the tensor's data and its gradient. No ``grad`` field is included
+    if the tensor does not have a gradient.
+
+    See Also
+    --------
+    mygrad.load
+
+    Examples
+    --------
+    >>> import mygrad as mg
+    >>> from tempfile import TemporaryFile
+    >>> outfile = TemporaryFile()
+    >>> x = mg.arange(10.0)
+    >>> mg.save(outfile, x)
+    >>> _ = outfile.seek(0) # Only needed here to simulate closing & reopening file
+    >>> mg.load(outfile)
+    Tensor([0., 1., 2., 3., 4., 5., 6., 7., 8., 9.])
+
+    An example of saving a tensor that has an associated gradient.
+
+    >>> (x * x).backward()
+    >>> x.grad
+    array([ 0.,  2.,  4.,  6.,  8., 10., 12., 14., 16., 18.])
+    >>> outfile = TemporaryFile()
+    >>> x = mg.arange(10.0)
+    >>> mg.save(outfile, x)
+    >>> _ = outfile.seek(0) # Only needed here to simulate closing & reopening file
+    >>> loaded = mg.load(outfile)
+    >>> loaded
+    Tensor([0., 1., 2., 3., 4., 5., 6., 7., 8., 9.])
+    >>> loaded.grad
+    array([ 0.,  2.,  4.,  6.,  8., 10., 12., 14., 16., 18.])
+    """
+    if not isinstance(tensor, tb.Tensor):
+        raise TypeError(
+            f"mygrad.save requires a Tensor-type object, got type {type(tensor)}"
+        )
+
+    if tensor.grad is not None:
+        np.savez(file, data=tensor.data, grad=tensor.grad)
+    else:
+        np.savez(file, data=tensor.data)
+
+
+def load(file: _FileLike) -> tb.Tensor:
+    """Loads a saved Tensor and its gradient information (if applicable).
+
+    This docstring was adapted from that of numpy.load()
+
+    Parameters
+    ----------
+    file : str | Path | BinaryIO
+        The name of the file that holds the tensor data to load.
+
+    Returns
+    -------
+    loaded : Tensor
+        The loaded tensor (whose gradient will be loaded if it was saved).
+
+    See Also
+    --------
+    mygrad.save
+
+    Examples
+    --------
+    >>> import mygrad as mg
+    >>> from tempfile import TemporaryFile
+    >>> outfile = TemporaryFile()
+    >>> x = mg.arange(10.0)
+    >>> mg.save(outfile, x)
+    >>> _ = outfile.seek(0) # Only needed here to simulate closing & reopening file
+    >>> mg.load(outfile)
+    Tensor([0., 1., 2., 3., 4., 5., 6., 7., 8., 9.])
+
+    An example of saving a tensor that has an associated gradient.
+
+    >>> (x * x).backward()
+    >>> x.grad
+    array([ 0.,  2.,  4.,  6.,  8., 10., 12., 14., 16., 18.])
+    >>> outfile = TemporaryFile()
+    >>> x = mg.arange(10.0)
+    >>> mg.save(outfile, x)
+    >>> _ = outfile.seek(0) # Only needed here to simulate closing & reopening file
+    >>> loaded = mg.load(outfile)
+    >>> loaded
+    Tensor([0., 1., 2., 3., 4., 5., 6., 7., 8., 9.])
+    >>> loaded.grad
+    array([ 0.,  2.,  4.,  6.,  8., 10., 12., 14., 16., 18.])
+    """
+    loaded = np.load(file)
+
+    loaded_tensor = tb.tensor(loaded["data"])
+
+    if "grad" in loaded:
+        loaded_tensor.backward(loaded["grad"])
+
+    return loaded_tensor


Test patch:
diff --git a/tests/conftest.py b/tests/conftest.py
index ef0cfced..abaa03a3 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -1,4 +1,5 @@
 import os
+import tempfile
 import warnings
 
 import pytest
@@ -70,3 +71,24 @@ def raise_on_mem_locking_state_leakage() -> bool:
         assert True
 
     clear_all_mem_locking_state()
+
+
+@pytest.fixture()
+def cleandir() -> str:
+    """This fixture will use the stdlib `tempfile` module to
+    move the current working directory to a tmp-dir for the
+    duration of the test.
+
+    Afterwards, the session returns to its previous working
+    directory, and the temporary directory and its contents
+    are removed.
+
+    Yields
+    ------
+    str
+        The name of the temporary directory."""
+    with tempfile.TemporaryDirectory() as tmpdirname:
+        old_dir = os.getcwd()
+        os.chdir(tmpdirname)
+        yield tmpdirname
+        os.chdir(old_dir)
diff --git a/tests/test_io.py b/tests/test_io.py
new file mode 100644
index 00000000..c5f54852
--- /dev/null
+++ b/tests/test_io.py
@@ -0,0 +1,50 @@
+from pathlib import Path
+from string import ascii_lowercase
+from typing import Any
+
+import hypothesis.strategies as st
+import pytest
+from hypothesis import given, settings
+from numpy.testing import assert_array_equal
+
+from mygrad import Tensor, load, save
+from tests.custom_strategies import everything_except, tensors
+
+filenames = st.text(ascii_lowercase, min_size=1).map(lambda x: x + ".npz")
+
+
+@given(
+    fname=filenames,
+    as_path=st.booleans(),
+    tensor=tensors(
+        include_grad=st.booleans(), dtype=st.sampled_from(["float32", "float64"])
+    ),
+)
+@pytest.mark.usefixtures("cleandir")
+def test_save_load_roundtrip(fname: str, as_path: bool, tensor: Tensor):
+    if as_path:
+        fname = Path(fname)
+
+    save(fname, tensor)
+    loaded = load(fname)
+    assert_array_equal(tensor, loaded)
+
+    if tensor.grad is None:
+        assert loaded.grad is None
+    else:
+        assert_array_equal(tensor.grad, loaded.grad)
+
+
+@settings(max_examples=10)
+@given(
+    fname=filenames,
+    as_path=st.booleans(),
+    not_tensor=everything_except(Tensor),
+)
+@pytest.mark.usefixtures("cleandir")
+def test_validation(fname: str, as_path: bool, not_tensor: Any):
+    if as_path:
+        fname = Path(fname)
+
+    with pytest.raises(TypeError):
+        save(fname, not_tensor)