| 27 | self.assertEqual(out.item(), i*2) |
| 28 | |
| 29 | def test_implicit_input(self): |
| 30 | # x is the implicit input (like a weight) |
| 31 | x = Tensor([0]) |
| 32 | |
| 33 | # this function has an implicit input and an explicit output |
| 34 | @TinyJit |
| 35 | def f(): |
| 36 | ret:Tensor = x*2 |
| 37 | return ret |
| 38 | |
| 39 | for i in range(5): |
| 40 | # NOTE: this must be realized here, otherwise the update doesn't happen |
| 41 | # if we were explicitly tracking the implicit input Tensors, we might not need this realize |
| 42 | x.assign(Tensor([i])).realize() |
| 43 | out = f() |
| 44 | self.assertEqual(out.item(), i*2) |
| 45 | |
| 46 | def test_implicit_output(self): |
| 47 | # out is the implicit output (it's assigned to) |