diff --git a/test/test_tensor_variable.py b/test/test_tensor_variable.py new file mode 100644 index 0000000000..b1613b9c41 --- /dev/null +++ b/test/test_tensor_variable.py @@ -0,0 +1,61 @@ +import unittest +import numpy as np +from tinygrad import Tensor, Variable + +class TestTensorVariable(unittest.TestCase): + def test_add_tvar(self): + vv = Variable("a", 0, 10) + vv.bind(1) + ret = (Tensor(vv) + 3).item() + assert ret == 4 + + def test_inner_tvar_node(self): + vv = Variable("w", 0, 10) + vv.bind(2) + ret = Tensor.from_node(vv * 4).item() + assert ret == 8 + + def test_inner_tvar_mul(self): + vv = Variable("w", 0, 10) + vv.bind(2) + assert (Tensor(3) * vv).item() == 6 + + def test_inner_tvar_mul_node(self): + vv = Variable("w", 0, 10) + vv.bind(2) + assert (Tensor(3) * (vv * 4)).item() == 24 + + def test_symbolic_mean(self): + vv = Variable("a", 1, 10) + vv.bind(2) + t = Tensor.ones(2, 2).contiguous().reshape(2, vv) + ret = t.mean().item() + assert ret == 1 + + def test_symbolic_mean_2d(self): + vv = Variable("a", 1, 10) + vv.bind(2) + vv2 = Variable("b", 1, 10) + vv2.bind(2) + t = Tensor.ones(2, 2).contiguous().reshape(vv2, vv) + ret = t.mean().item() + assert ret == 1 + + def test_symbolic_mean_2d_axis_1(self): + vv = Variable("a", 1, 10) + vv.bind(2) + vv2 = Variable("b", 1, 10) + vv2.bind(2) + t = Tensor.ones(2, 2).contiguous().reshape(vv2, vv) + ret = t.mean(axis=1).reshape(2, 1).numpy() + assert np.all(ret == 1) + + @unittest.skip("symbolic arange isn't supported") + def test_symbolic_arange(self): + vv = Variable("a", 1, 10) + vv.bind(2) + ret = Tensor.arange(0, vv) + ret.realize() + +if __name__ == '__main__': + unittest.main() diff --git a/test/unit/test_helpers.py b/test/unit/test_helpers.py index a0d0f41d45..bcdf779653 100644 --- a/test/unit/test_helpers.py +++ b/test/unit/test_helpers.py @@ -143,7 +143,7 @@ class TestRoundUp(unittest.TestCase): class TestFetch(unittest.TestCase): def test_fetch_bad_http(self): - self.assertRaises(Exception, fetch, 'http://www.google.com/404') + self.assertRaises(Exception, fetch, 'http://www.google.com/404', allow_caching=False) @unittest.skipIf(not CI, "pre commit tests should run offline") def test_fetch_small(self): diff --git a/tinygrad/codegen/linearizer.py b/tinygrad/codegen/linearizer.py index ba6530b2c3..ee12878c80 100644 --- a/tinygrad/codegen/linearizer.py +++ b/tinygrad/codegen/linearizer.py @@ -45,7 +45,8 @@ class Linearizer(Kernel): # NOTE: the consts have to be cached for deduping of downstream uops to work def const(self, b:ConstType, dtype:DType=dtypes.int32, insert_before:Optional[UOp|int]=None) -> UOp: - return self.uops.add(UOps.CONST, dtype, tuple(), b, insert_before=insert_before) + if isinstance(b, Variable): return self.uops.add(UOps.DEFINE_VAR, dtype, tuple(), b.unbind()[0], insert_before=insert_before) + else: return self.uops.add(UOps.CONST, dtype, tuple(), b, insert_before=insert_before) def cast(self, val: Tuple[UOp], dtype:DType, insert_before:Optional[UOp|int]=None) -> UOp: return self.uops.add(UOps.CAST, dtype, val, insert_before=insert_before) diff --git a/tinygrad/engine/schedule.py b/tinygrad/engine/schedule.py index 6b344eddd0..e4ca558ef7 100644 --- a/tinygrad/engine/schedule.py +++ b/tinygrad/engine/schedule.py @@ -38,6 +38,7 @@ def _recursive_lazyop(buf:LazyBuffer, inputs:List[LazyBuffer], outbufs:Tuple[Laz if buf.op is LoadOps.CONST: unbound_st, st_var_vals = st.simplify().unbind() var_vals.update(st_var_vals) + if isinstance(buf.arg, Variable): var_vals.__setitem__(*buf.arg.unbind()) return LazyOp(BufferOps.CONST, (), ConstBuffer(buf.arg, buf.dtype, unbound_st)) # if we aren't fusing it, it's a load and we add it to the inputs diff --git a/tinygrad/lazy.py b/tinygrad/lazy.py index 7002c63dd9..ab899f9c71 100644 --- a/tinygrad/lazy.py +++ b/tinygrad/lazy.py @@ -4,7 +4,7 @@ from typing import Union, Optional, Any, Tuple, List from tinygrad.dtype import dtypes, DType, ConstType, least_upper_dtype from tinygrad.helpers import prod, getenv, all_int, all_same, DEBUG from tinygrad.ops import LoadOps, UnaryOps, BinaryOps, TernaryOps, ReduceOps, Op, exec_alu, python_alu -from tinygrad.shape.symbolic import sint +from tinygrad.shape.symbolic import sint, Variable from tinygrad.shape.shapetracker import ShapeTracker from tinygrad.buffer import Buffer from weakref import ref, ReferenceType, WeakValueDictionary @@ -13,7 +13,7 @@ lazycache: WeakValueDictionary[Any, LazyBuffer] = WeakValueDictionary() def create_lazybuffer(device:str, st:ShapeTracker, dtype:DType, op:Optional[Op]=None, arg:Any=None, srcs:Tuple[LazyBuffer, ...]=(), base:Optional[LazyBuffer]=None, enable_cache=bool(getenv("LAZYCACHE", 1))): if st.size == 0: op, arg, srcs, base = LoadOps.CONST, 0, (), None - if op is LoadOps.CONST: arg, enable_cache = dtypes.as_const(arg, dtype), True + if op is LoadOps.CONST: arg, enable_cache = dtypes.as_const(arg, dtype) if not isinstance(arg, Variable) else arg, True cache_key = (device, st, dtype, op, arg, tuple(ref(x) for x in srcs)) if base is None else (st, ref(base)) if enable_cache and (rret := lazycache.get(cache_key, None)): return rret @@ -97,7 +97,7 @@ class LazyBuffer: new_shape = new_shape[:-1] + ((new_shape[-1]*self.dtype.itemsize) // dtype.itemsize,) return create_lazybuffer(self.device, ShapeTracker.from_shape(new_shape), dtype, UnaryOps.CAST, (dtype, bitcast), (self,)) - def is_unrealized_const(self): return self.base.realized is None and self.base.op is LoadOps.CONST + def is_unrealized_const(self): return self.base.realized is None and self.base.op is LoadOps.CONST and not isinstance(self.base.arg, Variable) def is_unrealized_unmasked_const(self): return self.is_unrealized_const() and all(v.mask is None for v in self.st.views) def _copy(self, device:str) -> LazyBuffer: diff --git a/tinygrad/ops.py b/tinygrad/ops.py index ce64b94d1e..7f3919cae7 100644 --- a/tinygrad/ops.py +++ b/tinygrad/ops.py @@ -87,7 +87,9 @@ class LazyOp: @functools.cached_property def lazyops(self) -> List[LazyOp]: return dedup([self] + [item for x in self.src for item in x.lazyops]) def vars(self) -> List[Variable]: - return sorted(set.union(*[x.arg.st.vars() for x in self.lazyops if x.op in BufferOps], set()), key=lambda x: str(x.expr)) + extract_vars = [x.arg.st.vars() for x in self.lazyops if x.op in BufferOps] + const_vars = [x.arg.val.unbind()[0] for x in self.lazyops if x.op is BufferOps.CONST and isinstance(x.arg.val, Variable)] + return sorted(set.union(*extract_vars, set(const_vars)), key=lambda x: str(x.expr)) def copy_ast(sz) -> LazyOp: rd = LazyOp(BufferOps.LOAD, (), MemBuffer(1, dtypes.uint8, st:=ShapeTracker.from_shape((sz,)))) diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index 4d6043ef9e..66be1f9b60 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -14,7 +14,7 @@ from tinygrad.features.multi import MultiLazyBuffer from tinygrad.ops import LoadOps, ScheduleItem from tinygrad.buffer import Buffer, BufferOptions from tinygrad.device import Device -from tinygrad.shape.symbolic import sint, Variable +from tinygrad.shape.symbolic import sint, Variable, MulNode, Node from tinygrad.engine.realize import run_schedule, memory_planner from tinygrad.engine.schedule import create_schedule_with_vars @@ -91,7 +91,7 @@ class Tensor: def __init__(self, mode:bool = True): self.mode = mode def __enter__(self): self.prev, Tensor.no_grad = Tensor.no_grad, self.mode def __exit__(self, exc_type, exc_value, traceback): Tensor.no_grad = self.prev - def __init__(self, data:Union[None, ConstType, List, Tuple, LazyBuffer, np.ndarray, bytes, MultiLazyBuffer], + def __init__(self, data:Union[None, ConstType, List, Tuple, LazyBuffer, np.ndarray, bytes, MultiLazyBuffer, Variable], device:Optional[Union[str, tuple, list]]=None, dtype:Optional[DType]=None, requires_grad:Optional[bool]=None): assert dtype is None or isinstance(dtype, DType), f"invalid dtype {dtype}" device = tuple(Device.canonicalize(x) for x in device) if isinstance(device, (tuple, list)) else Device.canonicalize(device) @@ -106,6 +106,7 @@ class Tensor: self._ctx: Optional[Function] = None if isinstance(data, LazyBuffer): assert dtype is None or dtype == data.dtype, "dtype doesn't match, and casting isn't supported" elif isinstance(data, get_args(ConstType)): data = _loadop(LoadOps.CONST, tuple(), dtype or dtypes.from_py(data), device, data) + elif isinstance(data, Variable): data = _loadop(LoadOps.CONST, tuple(), dtype or dtypes.from_py(data.unbind()[1]), device, data) elif isinstance(data, bytes): data = _fromcpu(np.frombuffer(data, np.uint8)) elif data is None: data = _loadop(LoadOps.EMPTY, (0,), dtype or dtypes.default_float, device) elif isinstance(data, list): @@ -267,6 +268,12 @@ class Tensor: self.lazydata = self.shard(devices, axis).lazydata return self + @staticmethod + def from_node(y:Node, **kwargs) -> Tensor: + if isinstance(y, MulNode): return Tensor.from_node(y.a, **kwargs) * y.b + if isinstance(y, Variable): return Tensor(y, **kwargs, requires_grad=False) + raise RuntimeError(f"unhandled Node {y}") + # ***** creation llop entrypoint ***** @staticmethod @@ -920,9 +927,8 @@ class Tensor: def min(self, axis=None, keepdim=False): return -((-self).max(axis=axis, keepdim=keepdim)) def mean(self, axis=None, keepdim=False): - assert all_int(self.shape), "does not support symbolic shape" out = self.sum(axis=axis, keepdim=keepdim, downcast_half=False) - return out.div(prod(self.shape) / prod(out.shape)).cast(self.dtype) if 0 not in out.shape else out.cast(self.dtype) + return out.div(prod(self.shape)).mul(prod(out.shape)).cast(self.dtype) if 0 not in out.shape else out.cast(self.dtype) def var(self, axis=None, keepdim=False, correction=1): assert all_int(self.shape), "does not support symbolic shape" square_sum = ((self - self.mean(axis=axis, keepdim=True)).square()).sum(axis=axis, keepdim=keepdim) @@ -1173,10 +1179,11 @@ class Tensor: x: Tensor = self if not isinstance(y, Tensor): # make y a Tensor - assert isinstance(y, (float, int, bool)), f"{type(y)=}, {y=}" + assert isinstance(y, (float, int, bool, Node)), f"{type(y)=}, {y=}" if isinstance(self.dtype, ImageDType) or dtypes.is_float(x.dtype) or (dtypes.is_int(x.dtype) and isinstance(y, int)): y_dtype = x.dtype else: y_dtype = dtypes.from_py(y) - y = Tensor(dtypes.as_const(y, y_dtype), self.device, y_dtype, requires_grad=False) + if isinstance(y, Node): y = Tensor.from_node(y, device=self.device) + else: y = Tensor(dtypes.as_const(y, y_dtype), self.device, y_dtype, requires_grad=False) if match_dtype: output_dtype = least_upper_dtype(x.dtype, y.dtype)