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.
64 lines
2.6 KiB
Python
64 lines
2.6 KiB
Python
import unittest
|
|
from tinygrad import Tensor
|
|
from tinygrad.uop.ops import PatternMatcher, Ops, UPat, graph_rewrite, RewriteContext, UOp
|
|
from tinygrad.engine.grouper import sym, merge_views
|
|
|
|
class TestRewriteTrackedChildren(unittest.TestCase):
|
|
@unittest.skip("track_children no longer supported")
|
|
def test_children_in_context(self):
|
|
def print_children(ctx:RewriteContext, sink:UOp):
|
|
view_w_child = sink.src[0].src[0].src[0]
|
|
assert view_w_child.op is Ops.VIEW
|
|
assert set([x.arg for x in ctx.children[view_w_child]]) == set((2,3))
|
|
ctx.update_children()
|
|
assert set([x.arg for x in ctx.children[view_w_child]]) == set((3,4))
|
|
# this is the 3
|
|
assert len(ctx.children[sink.src[0].src[1]]) == 1
|
|
assert next(iter(ctx.children[sink.src[0].src[1]])).op is Ops.ADD
|
|
# this is the 4
|
|
assert len(ctx.children[sink.src[0].src[0]]) == 1
|
|
assert next(iter(ctx.children[sink.src[0].src[0]])).op is Ops.ADD
|
|
rewrite = PatternMatcher([
|
|
(UPat(Ops.CONST, arg=2, name="x"), lambda x: x.replace(arg=4)),
|
|
(UPat(Ops.SINK, name="sink"), print_children)
|
|
])
|
|
a = Tensor(2)
|
|
b = Tensor(3)
|
|
c = a + b
|
|
sink = c.uop.sink()
|
|
sink = graph_rewrite(sink, rewrite, track_children=True)
|
|
|
|
def test_simple_child(self):
|
|
rewrite = PatternMatcher([
|
|
(UPat(Ops.CONST, arg=2, name="x"), lambda x: x.replace(arg=4)),
|
|
])
|
|
a = Tensor(2)
|
|
b = Tensor(3)
|
|
c = a + b
|
|
sink = c.uop
|
|
view_w_child = a.uop.src[0]
|
|
print([x().arg for x in view_w_child.children])
|
|
print([x.arg for x in sink.get_children_map()[view_w_child]])
|
|
self.assertSetEqual(set([x.arg for x in sink.get_children_map()[view_w_child]]), set((2,3)))
|
|
# children can either be added to or removed from the map with graph_rewrite
|
|
# added to is easy to detect, just hook the UOp constructor
|
|
# when are children removed?
|
|
# * if a rewrite rule returns a UOp, the matched node is removed from the graph
|
|
sink = graph_rewrite(sink, rewrite)
|
|
print([x().arg for x in view_w_child.children])
|
|
print([x.arg for x in sink.get_children_map()[view_w_child]])
|
|
self.assertSetEqual(set([x.arg for x in sink.get_children_map()[view_w_child]]), set((3,4)))
|
|
|
|
@unittest.skip("track_children no longer supported")
|
|
def test_child_after_parent_update(self):
|
|
def print_children(ctx, r):
|
|
ctx.update_children()
|
|
print(ctx.children[r])
|
|
extra = PatternMatcher([(UPat(Ops.REDUCE_AXIS, name="r"), print_children)])
|
|
a = Tensor.empty(3, 3)
|
|
r = (a+0).sum()
|
|
graph_rewrite(r.uop, merge_views+sym+extra, track_children=True)
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|