/
io.py
92 lines (73 loc) · 2.58 KB
/
io.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
import numpy as np
from pytensor.graph.basic import Apply, Constant
from pytensor.graph.op import Op
from pytensor.link.c.type import Generic
from pytensor.tensor.type import tensor
class LoadFromDisk(Op):
"""
An operation to load an array from disk.
See Also
--------
load
Notes
-----
Non-differentiable.
"""
__props__ = ("dtype", "shape", "mmap_mode")
def __init__(self, dtype, shape, mmap_mode=None):
self.dtype = np.dtype(dtype) # turn "float64" into np.float64
self.shape = shape
if mmap_mode not in (None, "c"):
raise ValueError(
"The only supported values for mmap_mode "
"are None and 'c', got %s" % mmap_mode
)
self.mmap_mode = mmap_mode
def make_node(self, path):
if isinstance(path, str):
path = Constant(Generic(), path)
return Apply(self, [path], [tensor(dtype=self.dtype, shape=self.shape)])
def perform(self, node, inp, out):
path = inp[0]
if path.split(".")[-1] == "npz":
raise ValueError(f"Expected a .npy file, got {path} instead")
result = np.load(path, mmap_mode=self.mmap_mode)
if result.dtype != self.dtype:
raise TypeError(
f"Expected an array of type {self.dtype}, got {result.dtype} instead"
)
out[0][0] = result
def __str__(self):
return (
f"Load{{dtype: {self.dtype}, shape: {self.shape}, mmep: {self.mmap_mode}}}"
)
def load(path, dtype, shape, mmap_mode=None):
"""
Load an array from an .npy file.
Parameters
----------
path
A Generic symbolic variable, that will contain a string
dtype : data-type
The data type of the array to be read.
shape
The static shape information of the loaded array.
mmap_mode
How the file will be loaded. None means that the
data will be copied into an array in memory, 'c' means that the file
will be mapped into virtual memory, so only the parts that are
needed will be actually read from disk and put into memory.
Other modes supported by numpy.load ('r', 'r+', 'w+') cannot
be supported by PyTensor.
Examples
--------
>>> from pytensor import *
>>> path = Variable(Generic(), None)
>>> x = tensor.load(path, 'int64', (None,))
>>> y = x*2
>>> fn = function([path], y)
>>> fn("stored-array.npy") # doctest: +SKIP
array([0, 2, 4, 6, 8], dtype=int64)
"""
return LoadFromDisk(dtype, shape, mmap_mode)(path)
__all__ = ["load"]