tensor variable (#4362)

* tensor variable support

* consttype without variable?

* __setitem__

* symbolic mean works

* arange test

* more tests

* a few more tests

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

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

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)

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

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:

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

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)