fix jit input_rawbuffers check wrt consts (#2689)

* fix jit input_rawbuffers check wrt consts

* .numpy()

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()

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