mirror of
https://github.com/dragonpilot/dragonpilot.git
synced 2026-06-23 23:12:04 +08:00
Vehicle Researcher
6adb63b915
date: 2026-06-04T09:49:56 master commit: c0ab3550eca2e9daf197c46b7e4b24aa9637cf2e
16 lines
525 B
Python
16 lines
525 B
Python
from tinygrad import Tensor, Device, TinyJit, dtypes
|
|
|
|
GPUS = Device[Device.DEFAULT].count()
|
|
N = 6144
|
|
|
|
@TinyJit
|
|
def many_matmul(A, B):
|
|
out = A
|
|
for _ in range(8): out = out@B
|
|
return out
|
|
|
|
if __name__ == "__main__":
|
|
A = Tensor.ones(GPUS, N, N, dtype=dtypes.half).shard(devices=tuple([f"{Device.DEFAULT}:{i}" for i in range(GPUS)]), axis=0).contiguous()
|
|
B = Tensor.ones(GPUS, N, N, dtype=dtypes.half).shard(devices=tuple([f"{Device.DEFAULT}:{i}" for i in range(GPUS)]), axis=0).contiguous()
|
|
while 1: many_matmul(A, B)
|