69 lines
2.4 KiB
69 lines
2.4 KiB
2 months ago
|
import random
|
||
|
|
from tinygrad.helpers import DEBUG, getenv
|
||
|
|
from test.unit.test_shapetracker import CheckingShapeTracker
|
||
|
|
|
||
|
|
def do_permute(st):
|
||
|
|
perm = list(range(0, len(st.shape)))
|
||
|
|
random.shuffle(perm)
|
||
|
|
perm = tuple(perm)
|
||
|
|
if DEBUG >= 1: print("st.permute(", perm, ")")
|
||
|
|
st.permute(perm)
|
||
|
|
|
||
|
|
def do_pad(st):
|
||
|
|
c = random.randint(0, len(st.shape)-1)
|
||
|
|
pad = tuple((random.randint(0,2), random.randint(0,2)) if i==c else (0,0) for i in range(len(st.shape)))
|
||
|
|
if DEBUG >= 1: print("st.pad(", pad, ")")
|
||
|
|
st.pad(pad)
|
||
|
|
|
||
|
|
def do_reshape_split_one(st):
|
||
|
|
c = random.randint(0, len(st.shape)-1)
|
||
|
|
poss = [n for n in [1,2,3,4,5] if st.shape[c]%n == 0]
|
||
|
|
spl = random.choice(poss)
|
||
|
|
shp = st.shape[0:c] + (st.shape[c]//spl, spl) + st.shape[c+1:]
|
||
|
|
if DEBUG >= 1: print("st.reshape(", shp, ")")
|
||
|
|
st.reshape(shp)
|
||
|
|
|
||
|
|
def do_reshape_combine_two(st):
|
||
|
|
if len(st.shape) < 2: return
|
||
|
|
c = random.randint(0, len(st.shape)-2)
|
||
|
|
shp = st.shape[:c] + (st.shape[c] * st.shape[c+1], ) + st.shape[c+2:]
|
||
|
|
if DEBUG >= 1: print("st.reshape(", shp, ")")
|
||
|
|
st.reshape(shp)
|
||
|
|
|
||
|
|
def do_shrink(st):
|
||
|
|
c = random.randint(0, len(st.shape)-1)
|
||
|
|
while 1:
|
||
|
|
shrink = tuple((random.randint(0,s), random.randint(0,s)) if i == c else (0,s) for i,s in enumerate(st.shape))
|
||
|
|
if all(x<y for (x,y) in shrink): break
|
||
|
|
if DEBUG >= 1: print("st.shrink(", shrink, ")")
|
||
|
|
st.shrink(shrink)
|
||
|
|
|
||
|
|
def do_stride(st):
|
||
|
|
c = random.randint(0, len(st.shape)-1)
|
||
|
|
stride = tuple(random.choice([-2,-1,2]) if i==c else 1 for i in range(len(st.shape)))
|
||
|
|
if DEBUG >= 1: print("st.stride(", stride, ")")
|
||
|
|
st.stride(stride)
|
||
|
|
|
||
|
|
def do_flip(st):
|
||
|
|
c = random.randint(0, len(st.shape)-1)
|
||
|
|
stride = tuple(-1 if i==c else 1 for i in range(len(st.shape)))
|
||
|
|
if DEBUG >= 1: print("st.stride(", stride, ")")
|
||
|
|
st.stride(stride)
|
||
|
|
|
||
|
|
def do_expand(st):
|
||
|
|
c = [i for i,s in enumerate(st.shape) if s==1]
|
||
|
|
if len(c) == 0: return
|
||
|
|
c = random.choice(c)
|
||
|
|
expand = tuple(random.choice([2,3,4]) if i==c else s for i,s in enumerate(st.shape))
|
||
|
|
if DEBUG >= 1: print("st.expand(", expand, ")")
|
||
|
|
st.expand(expand)
|
||
|
|
|
||
|
|
shapetracker_ops = [do_permute, do_pad, do_shrink, do_reshape_split_one, do_reshape_combine_two, do_stride, do_expand]
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
random.seed(42)
|
||
|
|
for _ in range(getenv("CNT", 200)):
|
||
|
|
st = CheckingShapeTracker((random.randint(2, 10), random.randint(2, 10), random.randint(2, 10)))
|
||
|
|
for i in range(8): random.choice(shapetracker_ops)(st)
|
||
|
|
st.assert_same()
|