Add Tensor.stack() and Tensor.repeat() (...trying to make einops work with tinygrad) (#758)

* add stack() and repeat() methods

* make stack a static method

.gitignore

@@ -1,4 +1,6 @@
 __pycache__
+.venv/
+.vscode
 notebooks
 .*.swp
 .*.swo

test/test_ops.py

@@ -575,6 +575,30 @@ class TestOps(unittest.TestCase):
     for dim in range(-1, 2):
       helper_test_op([(45,65), (45,65), (45,65)], lambda x,y,z: torch.cat((x,y,z), dim), lambda x,y,z: x.cat(y, z, dim=dim))
 
+  def test_stack(self):
+    x = Tensor.randn(45, 65, 3)
+
+    for dim in range(-1, 3):
+      helper_test_op([(45, 65, 3), (45, 65, 3), (45, 65, 3)], lambda x, y, z: torch.stack((x, y, z), dim=dim), lambda x, y, z: Tensor.stack([x, y, z], dim=dim))
+    
+    with self.assertRaises(IndexError):
+      Tensor.stack([x], dim=77)
+    
+  def test_repeat(self):
+    x = Tensor.randn(45, 65, 3)
+    base_repeats = [2, 4, 3]
+
+    for reps in [[], [4], [2, 1], [3, 2, 2]]:
+      repeats = base_repeats + reps
+      helper_test_op([(45, 65, 3)], lambda x: x.repeat(*repeats), lambda x: x.repeat(repeats))
+
+    with self.assertRaises(AssertionError):
+      x.repeat((2, 4))
+
+    with self.assertRaises(AssertionError):
+      x.repeat((2, 0, 4))
+       
+
   def test_clip(self):
     helper_test_op([(45,65)], lambda x: x.clip(-2.3, 1.2), lambda x: x.clip(-2.3, 1.2))
 

tinygrad/tensor.py

@@ -251,6 +251,23 @@ class Tensor:
       s[dim] = (-k, shape_cumsum[-1]-k)
     return functools.reduce(Tensor.__add__, [arg.slice(s) for arg,s in zip(catargs, slc)])
 
+  @staticmethod
+  def stack(tensors, dim=0):
+    first = tensors[0].unsqueeze(dim)
+    unsqueezed_tensors = [tensor.unsqueeze(dim) for tensor in tensors[1:]]
+    # checks for shapes and number of dimensions delegated to cat
+    return first.cat(*unsqueezed_tensors, dim=dim)
+
+  def repeat(self, repeats):
+    ndim = len(self.shape)
+    base_shape = self.shape
+    if len(repeats) > ndim:
+      base_shape = (1,) * (len(repeats) - ndim) + base_shape
+    new_shape = [x for i in range(len(base_shape)) for x in [1, base_shape[i]]]
+    expand_shape = [x for r,s in zip(repeats, base_shape) for x in [r,s]]
+    final_shape = [r*s for r,s in zip(repeats, base_shape)]
+    return self.reshape(new_shape).expand(expand_shape).reshape(final_shape)
+
   # TODO: make this nicer with syntactic sugar in slice
   def chunk(self, num, dim):
     slice_params = [[(0, s) for s in self.shape] for _ in range(num)]