mirror of
https://github.com/tinygrad/tinygrad.git
synced 2026-02-08 05:35:11 -05:00
c9493e41a6
* reorder expand * symbolic ops needs resolve here * s/arg/st + whitespace * viz --------- Co-authored-by: qazal <qazal.software@gmail.com>
13 lines
374 B
Python
13 lines
374 B
Python
from tinygrad import Tensor, GlobalCounters, Context
|
|
|
|
if __name__ == "__main__":
|
|
with Context(TRACK_MATCH_STATS=0): test = Tensor.ones(32, 10).contiguous().realize()
|
|
GlobalCounters.reset()
|
|
|
|
# this is the softmax from scaled_dot_product_attention
|
|
# it becomes 3 kernels
|
|
print("*** softmax ***")
|
|
with Context(NOOPT=1):
|
|
out = test.softmax(-1)
|
|
out.realize()
|