Tiny backend test_ops fix part 3 (tinygrad#9483)

* extract straightforward things from tinygrad#9302

* pass dtype and device for ones_like

Author: Anish9901

extra/torch_backend/backend.py

@@ -31,6 +31,7 @@ def device_count(self): return getenv("GPUS", 1) # TODO: device count in tiny?
 torch.utils.rename_privateuse1_backend("tiny")
 torch._register_device_module("tiny", TinyBackend())
 torch.utils.generate_methods_for_privateuse1_backend()
+aten = torch.ops.aten
 
 # in place operations with views
 def is_view(self: torch.Tensor) -> bool: return getattr(self, "_base", None) is not None
@@ -75,9 +76,37 @@ def _index_put_impl_(self, indices, values, accumulate=False, unsafe=False):
 def index_tensor(x, y):
   return aten.index(x.cpu(), [z.cpu() if isinstance(z, torch.Tensor) else None for z in y]).to(x.device)
 
+@torch.library.impl("aten::index_put", "privateuseone")
+def index_put(self, indices, values, accumulate=False):
+  return aten.index_put(self.cpu(), [z.cpu() if isinstance(z, torch.Tensor) else None for z in indices], values.cpu(), accumulate).tiny()
+
 @torch.library.impl("aten::randperm.generator_out", "privateuseone")
 def randperm_generator(n, generator=None, out=None): out.copy_(torch.randperm(n, generator=generator, device="cpu").tiny())
 
+@torch.library.impl("aten::cumprod", "privateuseone")
+# TODO: move to tinygrad
+def cumprod(self, dim, dtype=None): return aten.cumprod(self.cpu(), dim, dtype=dtype).tiny()
+
+@torch.library.impl("aten::cummax", "privateuseone")
+def cummax(self, dim):
+  # TODO: support cummax with indices to match torch
+  cummax, indices = aten.cummax(self.cpu(), dim)
+  return (cummax.tiny(), indices.tiny())
+
+@torch.library.impl("aten::nonzero", "privateuseone")
+# TODO: move to tinygrad
+def nonzero(self): return aten.nonzero(self.cpu()).tiny()
+
+def upsample_backward(grad_out, output_size, input_size, *args, f=None): return f(grad_out.cpu(), output_size, input_size, *args).tiny()
+
+for i in [
+  "upsample_linear1d_backward", "upsample_nearest1d_backward", "_upsample_nearest_exact1d_backward",
+  "upsample_nearest2d_backward", "_upsample_nearest_exact2d_backward",
+  "upsample_nearest3d_backward", "_upsample_nearest_exact3d_backward",
+  "upsample_trilinear3d_backward", "upsample_bilinear2d_backward"
+]:
+  torch.library.impl(f"aten::{i}", "privateuseone")(functools.partial(upsample_backward, f=getattr(aten, i)))
+
 # *** end bad functions on CPU ***
 
 @torch.library.impl("aten::zero_", "privateuseone")
@@ -162,24 +191,58 @@ def arange_start_step(start, end, step, dtype=None, device=None, pin_memory=None
 def convolution_overrideable(input, weight, bias, stride, padding, dilation, transposed, output_padding, groups):
   if TORCH_DEBUG >= 1:
     print(f"convolution {input.shape=} {weight.shape=} {stride=} {padding=} {dilation=} {transposed=} {output_padding=} {groups=}")
-  return wrap(unwrap(input).conv2d(unwrap(weight), unwrap(bias) if bias is not None else None,
-                                   groups=groups, stride=stride, dilation=dilation, padding=padding))
+  input, weight, bias = unwrap(input), unwrap(weight), unwrap(bias) if bias is not None else None
+  if not transposed: return wrap(input.conv2d(weight, bias, groups=groups, stride=stride, dilation=dilation, padding=padding))
+  return wrap(input.conv_transpose2d(weight, bias, groups=groups, stride=stride, dilation=dilation, padding=padding, output_padding=output_padding))
 
 @torch.library.impl("aten::convolution_backward_overrideable", "privateuseone")
 def convolution_backward_overrideable(grad_out, input, weight, stride, padding, dilation, transposed, output_padding, groups, output_mask):
   if TORCH_DEBUG >= 1:
     print(f"convolution_backward {input.shape=} {weight.shape=} {stride=} {padding=} {dilation=} {transposed=} {output_padding=} {groups=}")
   grad_out, input, weight, bias = unwrap(grad_out), unwrap(input), unwrap(weight), Tensor.zeros(weight.shape[0], device=_from_torch_device(weight.device))
-  out = Tensor.conv2d(input, weight, bias, groups=groups, stride=stride, dilation=dilation, padding=padding)
+  if not transposed: out = Tensor.conv2d(input, weight, bias, groups=groups, stride=stride, dilation=dilation, padding=padding)
+  else:
+    bias = Tensor.zeros(weight.shape[1] * groups)
+    out = Tensor.conv_transpose2d(input, weight, bias, groups=groups, stride=stride, dilation=dilation, padding=padding, output_padding=output_padding)
   grads = out.gradient(*[t for t,m in zip([input, weight, bias], output_mask) if m], gradient=grad_out)
   return tuple([wrap(grads.pop(0)) if m else None for m in output_mask])
 
+def avg_pool(self, kernel_size, stride=[], padding=0, ceil_mode=False, count_include_pad=True, divisor_override=None):
+  return wrap(unwrap(self).avg_pool2d(kernel_size, stride if stride != [] else None, padding=padding, ceil_mode=ceil_mode, count_include_pad=count_include_pad))
+
+def avg_pool_backward(grad_out, self, kernel_size, stride=None, padding=0, ceil_mode=False, count_include_pad=True, divisor_override=None):
+  self, grad_out = unwrap(self), unwrap(grad_out)
+  out = Tensor.avg_pool2d(self, kernel_size, stride if stride != [] else None, dilation=1, padding=padding, ceil_mode=ceil_mode, count_include_pad=count_include_pad)
+  return wrap(out.gradient(self, gradient=grad_out)[0])
+
+for dim in [2, 3]:
+  torch.library.impl(f"aten::avg_pool{dim}d", "privateuseone")(avg_pool)
+  torch.library.impl(f"aten::avg_pool{dim}d_backward", "privateuseone")(avg_pool_backward)
+
+def pad_forward(self, padding, mode=None): return wrap(Tensor.pad(unwrap(self), padding, mode=mode))
+
+def pad_backward(grad_out, self, padding, mode):
+  self, grad_out = unwrap(self), unwrap(grad_out)
+  out = Tensor.pad(self, padding, mode=mode)
+  return wrap(out.gradient(self, gradient=grad_out)[0])
+
+for dim in [1, 2, 3]:
+  for pad_type, mode in [("replication", "replicate"), ("reflection", "reflect")]:
+    torch.library.impl(f"aten::{pad_type}_pad{dim}d", "privateuseone")(functools.partial(pad_forward, mode=mode))
+    torch.library.impl(f"aten::{pad_type}_pad{dim}d_backward", "privateuseone")(functools.partial(pad_backward, mode=mode))
+
 def upsample(self, size, align_corners=False, mode=None): return wrap(Tensor.interpolate(unwrap(self), size, mode=mode, align_corners=align_corners))
 for i,pre in enumerate(["", "bi", "tri"]):
   torch.library.impl(f"aten::upsample_{pre}linear{i+1}d", "privateuseone")(functools.partial(upsample, mode="linear"))
   torch.library.impl(f"aten::upsample_nearest{i+1}d", "privateuseone")(functools.partial(upsample, mode="nearest"))
   torch.library.impl(f"aten::_upsample_nearest_exact{i+1}d", "privateuseone")(functools.partial(upsample, mode="nearest-exact"))
 
+@torch.library.impl("aten::scatter_add.out", "privateuseone")
+def scatter_add(self, dim, index, src, out):
+  self, index, src, out = unwrap(self), unwrap(index), unwrap(src), unwrap(out)
+  if self.shape == (): return wrap(out.assign(src))
+  return wrap(out.assign(Tensor.scatter_reduce(self, dim, index, src, reduce='sum')))
+
 @torch.library.impl("aten::_copy_from", "privateuseone")
 def _copy_from(src: torch.Tensor, dest, non_blocking=False):
   realize = dest.is_tiny and maybe_realize_storage(dest)
@@ -222,7 +285,6 @@ def sort_values(input, dim=-1, descending=False, stable=True, values=None, indic
 
 # register some decompositions
 from torch._decomp import get_decompositions
-aten = torch.ops.aten
 decomps = [
   aten.native_batch_norm, aten.native_batch_norm_backward,
   aten.native_layer_norm_backward,
@@ -344,7 +406,7 @@ def sort_values(input, dim=-1, descending=False, stable=True, values=None, indic
   "aten.scatter.value_out": Tensor.scatter,
   "aten.where.self_out": Tensor.where,
   "aten.prod.int_out": Tensor.prod,
-  "aten.scatter_add.out": functools.partial(Tensor.scatter_reduce, reduce='sum'),
+  "aten.scatter.src_out": Tensor.scatter,
   # NOTE: axis=[] in torch means all, change tinygrad?
   "aten.sum.IntList_out": lambda self,axis,keepdim=False,dtype=None:
     self.sum(axis if axis is None or len(axis) else None, keepdim,
@@ -408,9 +470,8 @@ def _wrap_out(*args, **kwargs):
   "aten.logical_not": Tensor.logical_not,
   "aten.logical_or_": inplace_fn("x")(lambda x, y: x.assign(x | y)),
   "aten.multinomial": Tensor.multinomial,
-  "aten.pad": Tensor.pad,
-  "aten.reflection_pad2d": functools.partial(Tensor.pad, mode="reflect"),
   "aten.masked_fill_.Scalar": inplace_fn("self")(lambda self, mask, value: self.assign(self.masked_fill(mask, value))),
+  "aten.masked_fill_.Tensor": inplace_fn("self")(lambda self, mask, value: self.assign(self.masked_fill(mask, value))),
   "aten.masked_fill.Scalar": Tensor.masked_fill,
   "aten.masked_fill.Tensor": Tensor.masked_fill,
   "aten.masked_select": Tensor.masked_select,
@@ -441,6 +502,9 @@ def _wrap_out(*args, **kwargs):
   "aten.repeat": Tensor.repeat,
   "aten.lerp.Tensor": Tensor.lerp,
   "aten.expand": Tensor.expand,
+  "aten.ones_like": lambda self, dtype=None, device=None, **kwargs:
+    self.ones_like(**{k: v for k, v in {"dtype": _from_torch_dtype(dtype) if dtype else None,
+                                        "device": _from_torch_device(device) if device else None}.items() if v is not None}),
   "aten.t": Tensor.transpose,
   "aten.detach": Tensor.detach,
   "aten.max.dim": lambda self, dim, keepdim=False: (self.max(dim, keepdim), self.argmax(dim, keepdim).cast(dtype=dtypes.int64))

test/test_ops.py

@@ -2671,6 +2671,7 @@ def test_gather(self):
                          vals=[[1., 2., 3.]])
 
   @unittest.expectedFailure
+  @unittest.skipIf(torch._C._get_privateuse1_backend_name() == "tiny", 'results in a success instead of a failure')
   def test_gather_failure(self):
     # gather with inf values do not work, other values results in nan
     helper_test_op(None, lambda x: x.gather(dim=0, index=torch.tensor([2, 1, 0, 1, 2], requires_grad=False)),