mirror of
https://github.com/tinygrad/tinygrad.git
synced 2026-01-10 07:28:15 -05:00
25 lines
573 B
Python
25 lines
573 B
Python
#!/usr/bin/env python3
|
|
import numpy as np
|
|
from tinygrad.tensor import Tensor
|
|
import time
|
|
|
|
# Tensor has max size of 0x4000 for now
|
|
ba = Tensor(np.random.normal(size=(0x4000,)))
|
|
for dev in ["CPU", "GPU", "ANE"]:
|
|
if dev == "GPU":
|
|
baa = ba.cuda()
|
|
elif dev == "ANE":
|
|
baa = ba.ane()
|
|
else:
|
|
baa = ba
|
|
for i in range(3):
|
|
st = time.time()
|
|
boaa = baa.relu()
|
|
et = time.time()
|
|
if i == 2:
|
|
print("%s can do at least %.2f MEGAReLUs/sec" % (dev, (np.prod(boaa.shape)/1e6)/(et-st)))
|
|
# decently reliable
|
|
assert(np.all(boaa.cpu().data >= 0))
|
|
|
|
|