more types

tinygrad/llops/ops_gpu.py

@@ -14,7 +14,7 @@ class CL:
   cl_ctx : Optional[cl.Context] = None
   cl_queue : Optional[cl.CommandQueue] = None
   def __init__(self):
-    if getattr(CL, "cl_queue", None) is not None: return
+    if CL.cl_queue is not None: return
     devices = sum([x.get_devices(device_type=cl.device_type.GPU) for x in cl.get_platforms()], [])
     if len(devices) == 0:  # settle for CPU
       devices = sum([x.get_devices(device_type=cl.device_type.CPU) for x in cl.get_platforms()], [])
@@ -91,7 +91,7 @@ class GPUBuffer:
     assert op == ProcessingOps.CONV, f"{op} isn't supported"
     return type(x)(C.out_shape)._processing_op([("input", x.contiguous_op()), ("weight", w.contiguous_op())], "acc", C)
 
-  def reduce_op(x, op:ReduceOps, new_shape:Tuple[int]):
+  def reduce_op(x, op:ReduceOps, new_shape:Tuple[int, ...]):
     ret = type(x)(new_shape)
     if op == ReduceOps.SUM: code, start = "out += a", "0.0"
     elif op == ReduceOps.MAX: code, start = "out = max(a,out)", "-INFINITY"

tinygrad/ops.py

@@ -163,7 +163,7 @@ class LazyBuffer:
   def binary_op(x:LazyBuffer, op:BinaryOps, y:LazyBuffer) -> LazyBuffer:
     return LazyBuffer(x.device, x.shape, BinaryOps, LazyOp(op, (x,y)))
 
-  def reduce_op(x:LazyBuffer, op:ReduceOps, new_shape:Tuple[int]) -> LazyBuffer:
+  def reduce_op(x:LazyBuffer, op:ReduceOps, new_shape:Tuple[int, ...]) -> LazyBuffer:
     return LazyBuffer(x.device, tuple(new_shape), ReduceOps, LazyOp(op, (x,), tuple(new_shape)))
 
   def movement_op(x:LazyBuffer, op:MovementOps, arg) -> LazyBuffer:

tinygrad/shapetracker.py

@@ -65,7 +65,7 @@ def strides_for_shape(shape):
   return tuple(strides)
 
 @functools.lru_cache(maxsize=None)
-def view_from_shape(shape:Tuple):
+def view_from_shape(shape:Tuple[int, ...]):
   if len(shape) == 0: shape = (1,)
   assert all([isinstance(x, int) for x in shape])
   return View(tuple(shape), strides_for_shape(shape))

tinygrad/tensor.py

@@ -1,8 +1,9 @@
 # inspired by https://github.com/karpathy/micrograd/blob/master/micrograd/engine.py
+from __future__ import annotations
 import inspect, functools, importlib
 import numpy as np
 from tinygrad.helpers import prod
-from typing import List
+from typing import List, Tuple, Callable, Optional
 from tinygrad.ops import Device
 
 from tinygrad.ops import LazyBuffer
@@ -178,17 +179,17 @@ class Tensor:
       ret += y.slice(arg=ts)
     return ret
 
-  def pad2d(self, padding):
+  def pad2d(self, padding:Tuple[int, ...]):
     # (padding_left, padding_right, padding_top, padding_bottom)
     return self[:, :, -padding[2]:self.shape[2]+padding[3], -padding[0]:self.shape[3]+padding[1]]
 
-  def matmul(x, w):
+  def matmul(x:Tensor, w:Tensor):
     # NOTE: we use a 1x1 conv2d to do the matmul. mxk @ kxn = (1,k,m,1).conv2d(n,k,1,1)
     bs, groups = prod(x.shape[0:-2]), prod(w.shape[0:-2])
     cin, cout = w.shape[-2], w.shape[-1]
     out_shape_t = tuple(list(x.shape[0:-2])+[cout,-1])
-    if len(x.shape) == 1: order, out_shape_t = (0,), (cout, )
-    else: order = tuple(list(range(len(x.shape)-2))+[len(x.shape)-1, len(x.shape)-2])
+    if len(x.shape) > 1: order = tuple(list(range(len(x.shape)-2))+[len(x.shape)-1, len(x.shape)-2])
+    else: order, out_shape_t = (0,), (cout, )
     worder = tuple(list(range(len(w.shape)-2))+[len(w.shape)-1, len(w.shape)-2])
 
     # NOTE: with NHWC we can remove the transposes
@@ -317,12 +318,12 @@ class Tensor:
   def reshape(self, shape): return self._reshape(shape=shape)
   def expand(self, shape): return self._expand(shape=shape)
 
-  def linear(self, weight, bias):
+  def linear(self, weight:Tensor, bias:Tensor):
     shp = [1] * (len(self.shape)-1) + [-1]
     ret = self.mul(weight.reshape(shape=shp)) if len(weight.shape) == 1 else self.dot(weight)
     return ret.add(bias.reshape(shape=shp))
 
-  def sequential(self, ll): return functools.reduce(lambda x,f: f(x), ll, self)
+  def sequential(self, ll:List[Callable[[Tensor], Tensor]]): return functools.reduce(lambda x,f: f(x), ll, self)
 
   def layernorm(x, eps=1e-5):
     y = (x - x.mean(axis=-1, keepdim=True))