9c7833faf9
* Vegetarian Filet o Fish model * fix.. atc.. * test cluster_speed_limit * fix.. cluster_speed_limit.. 2 * fix.. clusterspeedlimit3 * cruise speed to roadlimit speed * fix.. * fix.. eng * deltaUp/Down for lanechange * fix.. atc desire... * fix.. * ff * ff * fix.. * fix.. eng * fix engsound * Update desire_helper.py * fix.. connect... * fix curve_min speed * Revert "fix curve_min speed" This reverts commit fcc9c2eb14eb3504abef3e420db93e8882e56f37. * Reapply "fix curve_min speed" This reverts commit 2d2bba476c58a7b4e13bac3c3ad0e4694c95515d. * fix.. auto speed up.. roadlimit * fix.. atc auto lanechange... * Update desire_helper.py * Update cruise.py * debug atc... * fix.. waze alert offset.. * fix.. * test atc.. * fix.. * fix.. atc * atc test.. * fix.. atc * fix.. atc2 * fix.. atc3 * KerryGold Model. latsmooth_sec = 0.0 * lat smooth seconds 0.13 * fix comment * fix.. auto cruise, and speed unit * change lanemode switching. * erase mazda lkas button.
79 lines
2.6 KiB
Python
79 lines
2.6 KiB
Python
import unittest
|
|
from tinygrad import TinyJit, Tensor
|
|
|
|
# The JIT functions as a "capturing" JIT.
|
|
# Whatever kernels ran in the JIT the second run through the function will be the kernels that will run from then on.
|
|
# Explicit inputs to the function are updated in the JIT graph to the new inputs.
|
|
|
|
# JITs have four tensor types
|
|
# 1. Tensors that are explicit in the input, aka what's passed in. TODO: support lists/dicts/classes, anything get_state works on
|
|
# 2. Tensors that are explicit in the output, aka what's returned. TODO: same as above
|
|
# 3. Tensors that are implicit in the input as a closure.
|
|
# 4. Tensors that are implicit in the output because they were assigned to and realized.
|
|
|
|
# explicit inputs and outputs are realized on their way in and out of the JIT
|
|
# there's a whole bunch of edge cases and weirdness here that needs to be tested and clarified.
|
|
|
|
class TestJitCases(unittest.TestCase):
|
|
def test_explicit(self):
|
|
# this function has an explicit input and an explicit output
|
|
@TinyJit
|
|
def f(x:Tensor):
|
|
ret:Tensor = x*2
|
|
return ret
|
|
|
|
for i in range(5):
|
|
out = f(Tensor([i]))
|
|
self.assertEqual(out.item(), i*2)
|
|
|
|
def test_implicit_input(self):
|
|
# x is the implicit input (like a weight)
|
|
x = Tensor([0])
|
|
|
|
# this function has an implicit input and an explicit output
|
|
@TinyJit
|
|
def f():
|
|
ret:Tensor = x*2
|
|
return ret
|
|
|
|
for i in range(5):
|
|
# NOTE: this must be realized here, otherwise the update doesn't happen
|
|
# if we were explicitly tracking the implicit input Tensors, we might not need this realize
|
|
x.assign(Tensor([i])).realize()
|
|
out = f()
|
|
self.assertEqual(out.item(), i*2)
|
|
|
|
def test_implicit_output(self):
|
|
# out is the implicit output (it's assigned to)
|
|
out = Tensor([0])
|
|
|
|
# this function has an explicit input and an implicit output
|
|
@TinyJit
|
|
def f(x:Tensor):
|
|
# NOTE: this must be realized here
|
|
# if we were explicitly tracking the implicit output Tensors, we might not need this realize
|
|
out.assign(x*2).realize()
|
|
|
|
for i in range(5):
|
|
f(Tensor([i]))
|
|
self.assertEqual(out.item(), i*2)
|
|
|
|
def test_implicit_io(self):
|
|
# x is the implicit input (like a weight)
|
|
# out is the implicit output (it's assigned to)
|
|
x = Tensor([0])
|
|
out = Tensor([0])
|
|
|
|
# this function has an implicit input and an implicit output
|
|
@TinyJit
|
|
def f():
|
|
out.assign(x*2).realize() # NOTE: this must be realized here
|
|
|
|
for i in range(5):
|
|
x.assign(Tensor([i])).realize()
|
|
f()
|
|
self.assertEqual(out.item(), i*2)
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|