mirror of
https://github.com/tinygrad/tinygrad.git
synced 2026-02-19 02:44:40 -05:00
16 lines
701 B
Python
16 lines
701 B
Python
from tinygrad import Tensor, dtypes
|
|
dtypes.default_float = dtypes.float16
|
|
|
|
if __name__ == "__main__":
|
|
# matmuls in bert layers
|
|
BS = 96//6
|
|
tensors = [
|
|
(Tensor.empty(BS, 512, 1024), Tensor.empty(1024, 1024).T), # linear to get qkv
|
|
(Tensor.empty(BS, 512, 16, 64).permute(0,2,1,3), Tensor.empty(BS, 512, 16, 64).permute(0,2,3,1)), # q@k
|
|
(Tensor.empty(BS, 16, 512, 512), Tensor.empty(BS, 512, 16, 64).permute(0,2,1,3)), # qk@v
|
|
]
|
|
for t0, t1 in tensors:
|
|
print(f"{t0.shape=}, {t0.lazydata.st.real_strides()=}, {t1.shape=}, {t1.lazydata.st.real_strides()=}")
|
|
for _ in range(5):
|
|
t0.dot(t1, dtype="half").realize()
|