diff --git a/test/unit/test_function.py b/test/unit/test_function.py index 6a98e6ebef..67e0ca4d65 100644 --- a/test/unit/test_function.py +++ b/test/unit/test_function.py @@ -2,6 +2,7 @@ import numpy as np import unittest from tinygrad.function import function from tinygrad import Tensor +from tinygrad.uop.ops import UOp class TestFunction(unittest.TestCase): def test_simple(self): @@ -102,7 +103,6 @@ class TestFunction(unittest.TestCase): np.testing.assert_allclose(w.grad.numpy(), [4., 5., 6.]) def test_symbolic_index(self): - from tinygrad.uop.ops import UOp table = Tensor([10,20,30,40]).contiguous().realize() @function def f(x:Tensor, start_pos:int|UOp) -> Tensor: @@ -111,6 +111,14 @@ class TestFunction(unittest.TestCase): v = UOp.variable("start_pos", 0, 3) np.testing.assert_equal(f(Tensor([1,2,3]), v.bind(0)).numpy(), [11,12,13]) + def test_symbolic_shape_input(self): + table = Tensor([10,20,30,40]).contiguous().realize() + @function + def f(x:Tensor) -> Tensor: return x * 2 + sz = UOp.variable("sz", 1, 3) + slic = table[:sz.bind(2)] + np.testing.assert_equal(f(slic)[:2].numpy(), [20,40]) + def test_nested_calls(self): w = Tensor([10., 20., 30.]) @function diff --git a/tinygrad/function.py b/tinygrad/function.py index e67ff2382d..27b5ebaddb 100644 --- a/tinygrad/function.py +++ b/tinygrad/function.py @@ -10,7 +10,7 @@ def add_to_ctx(ctx, x:UOp): return ret pm_ctx = PatternMatcher([ - (UPat(Ops.BUFFER, name="x"), add_to_ctx), + (UPat((Ops.BUFFER, Ops.BIND), name="x"), add_to_ctx), (UPat((Ops.ASSIGN, Ops.CONTIGUOUS), name="x"), lambda ctx,x: add_to_ctx(ctx,x) if not x.op_in_backward_slice_with_self(Ops.PARAM) else None), ]) diff --git a/tinygrad/schedule/rangeify.py b/tinygrad/schedule/rangeify.py index 26942daef8..bc54519eff 100644 --- a/tinygrad/schedule/rangeify.py +++ b/tinygrad/schedule/rangeify.py @@ -90,7 +90,7 @@ def resolve_call(c:UOp, allow_param_mismatch=True) -> UOp|None: dict_map = {x:args[x.arg] for x in params} for i, (p, a) in enumerate(dict_map.items()): - if p.shape != a.shape: raise TypeError(f"arg {i} shape mismatch: expected {p.shape}, got {a.shape}") + if p.max_shape != a.max_shape: raise TypeError(f"arg {i} shape mismatch: expected {p.shape}, got {a.shape}") if p.dtype != a.dtype: raise TypeError(f"arg {i} dtype mismatch: expected {p.dtype}, got {a.dtype}") return c.src[0].substitute(dict_map, walk=True)