mirror of
https://github.com/tinygrad/tinygrad.git
synced 2026-01-06 21:53:53 -05:00
add test for jit footguns (#13701)
* add test for jit footguns * shorter * notes
This commit is contained in:
George Hotz
270
test/test_jit_footguns.py
Normal file
270
test/test_jit_footguns.py
Normal file
@@ -0,0 +1,270 @@
|
||||
#!/usr/bin/env python
|
||||
"""
|
||||
JIT Footguns: Documenting unexpected behavior changes when using @TinyJit
|
||||
|
||||
Each test shows behavior that works without JIT but changes with JIT.
|
||||
Comments marked "should be X!" indicate the intuitively expected value.
|
||||
|
||||
SILENT MISMATCHES (highest priority - wrong results, no error):
|
||||
tensors_in_containers_ignored EASY only checks t.__class__ is Tensor, could scan lists/dicts
|
||||
non_tensor_outputs_frozen EASY could warn/error if return contains non-Tensor values
|
||||
class_method_shared_across_instances EASY could check if first arg is self and warn
|
||||
output_buffer_reuse MED performance tradeoff, could add option or better docs
|
||||
python_constants_frozen HARD inherent to tracing JITs
|
||||
conditional_branches_frozen HARD inherent to tracing JITs
|
||||
|
||||
ERRORS RAISED (lower priority - at least users know):
|
||||
positional_kwargs_cannot_mix EASY normalize positional args to kwargs using function signature
|
||||
duplicate_inputs_fail MED would need to handle aliasing in input_replace
|
||||
nested_jit_fails_on_second_call MED could fail on first call instead of second
|
||||
"""
|
||||
import unittest
|
||||
import numpy as np
|
||||
from tinygrad import Tensor, TinyJit
|
||||
|
||||
class TestJitFootguns(unittest.TestCase):
|
||||
|
||||
def test_output_buffer_reuse(self):
|
||||
"""Output tensors share buffer after capture - old references get overwritten."""
|
||||
@TinyJit
|
||||
def f(x): return x.sum().realize()
|
||||
|
||||
r1 = f(Tensor([1, 1])) # warmup
|
||||
r2 = f(Tensor([2, 2])) # capture
|
||||
r3 = f(Tensor([3, 3])) # jit exec
|
||||
|
||||
self.assertEqual(r1.item(), 2) # warmup result independent
|
||||
self.assertEqual(r3.item(), 6) # latest is correct
|
||||
self.assertEqual(r2.item(), 6) # should be 4! (overwritten by r3)
|
||||
|
||||
def test_output_buffer_workaround(self):
|
||||
"""Use .clone().realize() to get independent copies."""
|
||||
@TinyJit
|
||||
def f(x): return x.sum().realize()
|
||||
|
||||
r1 = f(Tensor([1, 1])).clone().realize()
|
||||
r2 = f(Tensor([2, 2])).clone().realize()
|
||||
r3 = f(Tensor([3, 3])).clone().realize()
|
||||
|
||||
self.assertEqual([r1.item(), r2.item(), r3.item()], [2, 4, 6])
|
||||
|
||||
def test_non_tensor_outputs_frozen(self):
|
||||
"""Non-tensor return values are frozen at capture time."""
|
||||
@TinyJit
|
||||
def f(x, mult): return (x * 2).realize(), mult * 10
|
||||
|
||||
# collect results, copying tensor values immediately (buffer reuse!)
|
||||
results = []
|
||||
for i in range(5):
|
||||
t, s = f(Tensor([i]), i)
|
||||
results.append((t.item(), s))
|
||||
|
||||
# tensor outputs work correctly
|
||||
self.assertEqual([r[0] for r in results[2:]], [4, 6, 8])
|
||||
# scalar outputs frozen at capture (i=1) - should be 20, 30, 40!
|
||||
self.assertEqual([r[1] for r in results[2:]], [10, 10, 10])
|
||||
|
||||
def test_duplicate_inputs_fail(self):
|
||||
"""JIT cannot handle the same tensor passed as multiple arguments."""
|
||||
@TinyJit
|
||||
def f(a, b): return (a + b).realize()
|
||||
|
||||
x = Tensor([1, 2, 3])
|
||||
with self.assertRaises(AssertionError):
|
||||
f(x, x)
|
||||
|
||||
def test_tensors_in_containers_ignored(self):
|
||||
"""Tensors inside lists/dicts are not tracked as inputs."""
|
||||
@TinyJit
|
||||
def f(a, arr): return (a + arr[0]).realize()
|
||||
|
||||
results = []
|
||||
for i in range(4):
|
||||
a, b = Tensor([1, 1, 1]).realize(), Tensor([i, i, i]).realize()
|
||||
results.append(f(a, [b]).numpy().copy())
|
||||
|
||||
np.testing.assert_array_equal(results[0], [1, 1, 1]) # warmup
|
||||
np.testing.assert_array_equal(results[1], [2, 2, 2]) # capture
|
||||
np.testing.assert_array_equal(results[2], [2, 2, 2]) # should be [3,3,3]!
|
||||
np.testing.assert_array_equal(results[3], [2, 2, 2]) # should be [4,4,4]!
|
||||
|
||||
def test_nested_jit_fails_on_second_call(self):
|
||||
"""Nested JIT works on first call but fails on second."""
|
||||
@TinyJit
|
||||
def inner(t): return t + 1
|
||||
@TinyJit
|
||||
def outer(t): return inner(t) * 3
|
||||
|
||||
self.assertEqual(outer(Tensor([1])).realize().item(), 6) # works!
|
||||
with self.assertRaises(RuntimeError):
|
||||
outer(Tensor([2])).realize() # fails
|
||||
|
||||
def test_implicit_inputs_need_realize(self):
|
||||
"""Closure tensors must be realized before JIT call."""
|
||||
x = Tensor([0])
|
||||
|
||||
@TinyJit
|
||||
def f(): return (x * 2).realize()
|
||||
|
||||
for i in range(5):
|
||||
x.assign(Tensor([i])).realize() # must realize!
|
||||
self.assertEqual(f().item(), i * 2)
|
||||
|
||||
def test_views_with_different_offsets_fail(self):
|
||||
"""JIT requires consistent tensor views across calls."""
|
||||
@TinyJit
|
||||
def f(a): return (a + 1).realize()
|
||||
|
||||
base = Tensor.randn(10, 10).realize()
|
||||
with self.assertRaises(AssertionError):
|
||||
for i in range(1, 5):
|
||||
f(base[:, i:i+2]) # different offset each time
|
||||
|
||||
def test_shape_change_after_capture_fails(self):
|
||||
"""Shapes are locked at capture time."""
|
||||
@TinyJit
|
||||
def f(a, b): return (a + b).realize()
|
||||
|
||||
f(Tensor.randn(10, 10), Tensor.randn(10, 10)) # warmup
|
||||
f(Tensor.randn(10, 10), Tensor.randn(10, 10)) # capture
|
||||
|
||||
with self.assertRaises(AssertionError):
|
||||
f(Tensor.randn(20, 20), Tensor.randn(20, 20))
|
||||
|
||||
def test_python_constants_frozen(self):
|
||||
"""Python variables inside JIT use capture-time values."""
|
||||
mult = 1
|
||||
|
||||
@TinyJit
|
||||
def f(x): return (x * mult).realize()
|
||||
|
||||
results = []
|
||||
for i in range(5):
|
||||
mult = i + 1
|
||||
results.append(f(Tensor([10])).item())
|
||||
|
||||
self.assertEqual(results[0], 10) # warmup, mult=1
|
||||
self.assertEqual(results[1], 20) # capture, mult=2
|
||||
self.assertEqual(results[2], 20) # should be 30!
|
||||
self.assertEqual(results[3], 20) # should be 40!
|
||||
|
||||
def test_conditional_branches_frozen(self):
|
||||
"""Only the branch taken during capture runs thereafter."""
|
||||
@TinyJit
|
||||
def f(x, use_square):
|
||||
if use_square:
|
||||
return (x * x).realize()
|
||||
return (x * 2).realize()
|
||||
|
||||
f(Tensor([3]), True) # warmup
|
||||
f(Tensor([3]), False) # capture (False branch)
|
||||
|
||||
result = f(Tensor([3]), True) # passing True but False branch runs
|
||||
self.assertEqual(result.item(), 6) # should be 9!
|
||||
|
||||
def test_positional_kwargs_cannot_mix(self):
|
||||
"""Must use same calling convention after capture."""
|
||||
@TinyJit
|
||||
def f(a, b): return (a + b).realize()
|
||||
|
||||
f(Tensor([1]), Tensor([2])) # warmup with positional
|
||||
f(Tensor([1]), Tensor([2])) # capture with positional
|
||||
|
||||
with self.assertRaises(AssertionError):
|
||||
f(a=Tensor([3]), b=Tensor([4])) # kwargs fail
|
||||
|
||||
def test_class_method_shared_across_instances(self):
|
||||
"""JIT on instance methods is shared at class level."""
|
||||
class Model:
|
||||
def __init__(self, scale):
|
||||
self.scale = Tensor([scale])
|
||||
@TinyJit
|
||||
def forward(self, x):
|
||||
return (x * self.scale).realize()
|
||||
|
||||
m1, m2 = Model(2), Model(3)
|
||||
|
||||
m1.forward(Tensor([5])) # warmup
|
||||
m1.forward(Tensor([5])) # capture with m1.scale=2
|
||||
|
||||
self.assertEqual(m1.forward(Tensor([5])).item(), 10)
|
||||
self.assertEqual(m2.forward(Tensor([5])).item(), 10) # should be 15!
|
||||
|
||||
def test_side_effects_only_during_capture(self):
|
||||
"""Function body not executed during JIT replay."""
|
||||
call_count = [0]
|
||||
|
||||
@TinyJit
|
||||
def f(x):
|
||||
call_count[0] += 1
|
||||
return (x * 2).realize()
|
||||
|
||||
f(Tensor([1])) # warmup
|
||||
f(Tensor([2])) # capture
|
||||
self.assertEqual(call_count[0], 2)
|
||||
|
||||
f(Tensor([3]))
|
||||
f(Tensor([4]))
|
||||
f(Tensor([5]))
|
||||
self.assertEqual(call_count[0], 2) # still 2, not 5!
|
||||
|
||||
def test_nothing_realized_fails(self):
|
||||
"""Must JIT at least one kernel."""
|
||||
@TinyJit
|
||||
def f(a, b): return None
|
||||
|
||||
with self.assertRaises(AssertionError):
|
||||
for _ in range(3):
|
||||
f(Tensor([1]), Tensor([2]))
|
||||
|
||||
|
||||
class TestJitCorrectBehavior(unittest.TestCase):
|
||||
"""Behaviors that work correctly - documented for clarity."""
|
||||
|
||||
def test_random_regenerates(self):
|
||||
"""Random tensors regenerate each call."""
|
||||
@TinyJit
|
||||
def f(x):
|
||||
return (x + Tensor.rand(3)).realize()
|
||||
|
||||
f(Tensor([0, 0, 0])) # warmup
|
||||
f(Tensor([0, 0, 0])) # capture
|
||||
|
||||
results = {tuple(f(Tensor([0, 0, 0])).numpy().tolist()) for _ in range(5)}
|
||||
self.assertEqual(len(results), 5)
|
||||
|
||||
def test_unrealized_return_auto_realized(self):
|
||||
"""Unrealized return tensors are auto-realized."""
|
||||
@TinyJit
|
||||
def f(a, b): return a + b # no explicit realize
|
||||
|
||||
for _ in range(5):
|
||||
a, b = Tensor.randn(10), Tensor.randn(10)
|
||||
np.testing.assert_allclose(f(a, b).numpy(), a.numpy() + b.numpy(), atol=1e-5)
|
||||
|
||||
def test_kwargs_order_doesnt_matter(self):
|
||||
"""Kwargs are sorted by name, so order doesn't matter."""
|
||||
@TinyJit
|
||||
def f(first, second): return (first / second).realize()
|
||||
|
||||
for _ in range(3):
|
||||
a, b = Tensor.randn(10), Tensor.randn(10) + 1
|
||||
np.testing.assert_allclose(f(second=b, first=a).numpy(), a.numpy() / b.numpy(), atol=1e-4)
|
||||
np.testing.assert_allclose(f(first=a, second=b).numpy(), a.numpy() / b.numpy(), atol=1e-4)
|
||||
|
||||
def test_input_mutation_consistent(self):
|
||||
"""Input mutation via assign works consistently."""
|
||||
@TinyJit
|
||||
def f(x):
|
||||
x += 1
|
||||
x.realize()
|
||||
return x
|
||||
|
||||
a = Tensor([0]).contiguous().realize()
|
||||
for _ in range(5):
|
||||
f(a)
|
||||
self.assertEqual(a.item(), 5)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user