forked from minitorch/Module-3
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathtensor_strategies.py
More file actions
126 lines (105 loc) · 3.4 KB
/
Copy pathtensor_strategies.py
File metadata and controls
126 lines (105 loc) · 3.4 KB
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
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
from typing import List, Optional
from hypothesis import settings
from hypothesis.strategies import (
DrawFn,
SearchStrategy,
composite,
floats,
integers,
lists,
permutations,
)
import minitorch
from minitorch import Tensor, TensorBackend, TensorData, UserIndex, UserShape
from .strategies import small_ints
settings.register_profile("ci", deadline=None)
settings.load_profile("ci")
@composite
def vals(draw: DrawFn, size: int, number: SearchStrategy[float]) -> Tensor:
pts = draw(
lists(
number,
min_size=size,
max_size=size,
)
)
return minitorch.tensor(pts)
@composite
def shapes(draw: DrawFn) -> minitorch.UserShape:
lsize = draw(lists(small_ints, min_size=1, max_size=4))
return tuple(lsize)
@composite
def tensor_data(
draw: DrawFn,
numbers: SearchStrategy[float] = floats(),
shape: Optional[UserShape] = None,
) -> TensorData:
if shape is None:
shape = draw(shapes())
size = int(minitorch.prod(shape))
data = draw(lists(numbers, min_size=size, max_size=size))
permute: List[int] = draw(permutations(range(len(shape))))
permute_shape = tuple([shape[i] for i in permute])
z = sorted(enumerate(permute), key=lambda a: a[1])
reverse_permute = [a[0] for a in z]
td = minitorch.TensorData(data, permute_shape)
ret = td.permute(*reverse_permute)
assert ret.shape[0] == shape[0]
return ret
@composite
def indices(draw: DrawFn, layout: Tensor) -> UserIndex:
return tuple((draw(integers(min_value=0, max_value=s - 1)) for s in layout.shape))
@composite
def tensors(
draw: DrawFn,
numbers: SearchStrategy[float] = floats(
allow_nan=False, min_value=-100, max_value=100
),
backend: Optional[TensorBackend] = None,
shape: Optional[UserShape] = None,
) -> Tensor:
backend = minitorch.SimpleBackend if backend is None else backend
td = draw(tensor_data(numbers, shape=shape))
return minitorch.Tensor(td, backend=backend)
@composite
def shaped_tensors(
draw: DrawFn,
n: int,
numbers: SearchStrategy[float] = floats(
allow_nan=False, min_value=-100, max_value=100
),
backend: Optional[TensorBackend] = None,
) -> List[Tensor]:
backend = minitorch.SimpleBackend if backend is None else backend
td = draw(tensor_data(numbers))
values = []
for i in range(n):
data = draw(lists(numbers, min_size=td.size, max_size=td.size))
values.append(
minitorch.Tensor(
minitorch.TensorData(data, td.shape, td.strides), backend=backend
)
)
return values
@composite
def matmul_tensors(
draw: DrawFn,
numbers: SearchStrategy[float] = floats(
allow_nan=False, min_value=-100, max_value=100
),
) -> List[Tensor]:
i, j, k = [draw(integers(min_value=1, max_value=10)) for _ in range(3)]
l1 = (i, j)
l2 = (j, k)
values = []
for shape in [l1, l2]:
size = int(minitorch.prod(shape))
data = draw(lists(numbers, min_size=size, max_size=size))
values.append(minitorch.Tensor(minitorch.TensorData(data, shape)))
return values
def assert_close_tensor(a: Tensor, b: Tensor) -> None:
if a.is_close(b).all().item() != 1.0:
assert False, (
"Tensors are not close \n x.shape=%s \n x=%s \n y.shape=%s \n y=%s \n Diff=%s %s"
% (a.shape, a, b.shape, b, a - b, a.is_close(b))
)