diff --git a/test/test_jit.py b/test/test_jit.py index 8cacc59405..0a487b4037 100644 --- a/test/test_jit.py +++ b/test/test_jit.py @@ -286,6 +286,17 @@ class TestJit(unittest.TestCase): assert isinstance(jf.jit_cache[0].prg, graph_t) assert isinstance(jf.jit_cache[1].prg, graph_t) + def test_jit_const_inputs(self): + @TinyJit + def f(x,y): return (x+y).realize() + for _ in range(5): + np.testing.assert_equal(f(Tensor.ones(3), Tensor.zeros(3)).numpy(), np.ones(3)) + + @TinyJit + def g(x,y,z): return (x+y+z).realize() + for i in range(5): + np.testing.assert_equal(g(Tensor([i]*3), Tensor.ones(3), Tensor.zeros(3)).numpy(), np.array([i+1]*3)) + if __name__ == '__main__': unittest.main() diff --git a/tinygrad/jit.py b/tinygrad/jit.py index ceddd12e40..94a2739767 100644 --- a/tinygrad/jit.py +++ b/tinygrad/jit.py @@ -48,12 +48,13 @@ class TinyJit(Generic[ReturnType]): def __get__(self, obj, objtype): return functools.partial(self.__call__, obj) def __call__(self, *args, **kwargs) -> ReturnType: - # all inputs are realized + # all inputs (except const) are realized input_tensors: Dict[Union[int, str], Tensor] = {cast(Union[int, str], k):v.realize() for k,v in itertools.chain(enumerate(args), kwargs.items()) if v.__class__ is Tensor} expected_name_sts_dtype = tuple([(k, v.lazydata.st.unbind(), v.dtype) for k,v in input_tensors.items()]) # get rawbuffers - input_rawbuffers: List[Buffer] = [cast(Buffer, v.lazydata.realized) for v in input_tensors.values()] + # TODO: why can .realized have Any type? + input_rawbuffers: List[Buffer] = [cast(Buffer, v.lazydata.realized) for v in input_tensors.values() if v.lazydata.realized is not None] assert len(set(input_rawbuffers)) == len(input_rawbuffers), "duplicate inputs to JIT" # get variables: they can either be in Tensors or passed in as arguments, and all must be bound. these are all global