nimlgen/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Activity

more torch backend fixups (#17562)

Browse Source
This commit is contained in:
chenyu
2026-08-16 21:31:53 -04:00
committed by GitHub
parent 138fb4a783
commit e25f86721d
2 changed files with 53 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files
extra/torch_backend/backend.py
+15 -33
View File
@@ -4,7 +4,7 @@
# A006 Lambda argument `input` is shadowing a Python builtin
from tinygrad import Tensor, dtypes, Device
from tinygrad.uop.ops import Ops, GroupOp
from tinygrad.helpers import getenv, prod, strides_for_shape, argfix
from tinygrad.helpers import getenv, prod, strides_for_shape
import torch.lib
TORCH_DEBUG = getenv("TORCH_DEBUG")
import torch, pathlib, operator, functools, weakref
@@ -99,38 +99,11 @@ def _apply_view_ops(target, ops):
for fn, args, kwargs in ops: target = fn(target, *args, **kwargs)
return target
# similar to https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/InferSize.h
def _reshape_target_shape(shape:tuple[int, ...], args) -> tuple[int, ...]|None:
if not (req := argfix(*args)): return None
new_shape, infer_idx = [], -1
for i, s in enumerate(req):
if s is None: s = shape[i] if i < len(shape) else None
if not isinstance(s, int): return None
if s == -1:
if infer_idx != -1: return None
infer_idx = len(new_shape)
new_shape.append(s)
total = prod(shape)
if infer_idx != -1:
known = prod(x for x in new_shape if x != -1)
if known == 0:
if total != 0: return None
new_shape[infer_idx] = 0
else: new_shape[infer_idx] = total // known
return tuple(new_shape) if prod(new_shape) == total else None
# TODO: can we get rid of this? only for test_flatten_reshape_add
# a chain of reshapes (and detaches, which move nothing) is undone by reshaping the value back to the base
def _try_simple_reshape_view_write(base: Tensor, view: Tensor, val: Tensor) -> bool:
if not (ops := _get_view_ops(view)): return False
shapes = [base.shape]
for fn, args, _ in ops:
if fn is Tensor.detach: continue # detach leaves every element where it was (a tracked view only on torch<2.10)
if fn is not Tensor.reshape: return False
if not (next_shape := _reshape_target_shape(shapes[-1], args)): return False
shapes.append(next_shape)
if shapes[-1] != view.shape: return False
for s in reversed(shapes[:-1]): val = val.reshape(s)
base.assign(val)
if any(fn not in (Tensor.reshape, Tensor.detach) for fn, _, _ in ops): return False
base.assign(val.reshape(base.shape))
return True
def _view_write(base: Tensor, view: Tensor, value: Tensor) -> None:
@@ -445,6 +418,7 @@ def _linalg_svd(self, full_matrices=False):
from torch._decomp import get_decompositions
decomps = [
aten.native_layer_norm_backward,
aten.native_group_norm_backward,
aten.linalg_cross,
aten.addmm,
aten.addcmul,
@@ -479,7 +453,13 @@ decomps = [
aten._softmax_backward_data, aten.embedding_dense_backward,
aten.linalg_vector_norm,
aten.binary_cross_entropy, aten.binary_cross_entropy_backward,
# the C++ mse/smooth_l1 kernels resize their out tensor, and a tiny tensor has no storage to resize
aten.mse_loss, aten.mse_loss_backward,
aten.smooth_l1_loss, aten.smooth_l1_loss_backward,
aten.upsample_nearest2d.out,
# NOTE: only the "out" overload, the "vec" one is CompositeImplicitAutograd and overriding it loses the autograd kernel
aten.upsample_bicubic2d.out,
aten._adaptive_avg_pool2d,
# activations
aten.hardswish, aten.hardswish_backward,
aten.hardtanh, aten.hardtanh_backward,
@@ -661,8 +641,9 @@ tiny_backend = {**{k:wrap_out(v) for k,v in tiny_backend_out.items()}, **{
"aten.acos": Tensor.acos,
"aten.any": Tensor.any,
"aten.bitwise_not": Tensor.bitwise_not,
"aten.argmax": Tensor.argmax,
"aten.argmin": Tensor.argmin,
# tinygrad indexes with int32, torch's arg reduces return int64
"aten.argmax": lambda self, dim=None, keepdim=False: self.argmax(dim, keepdim).cast(dtypes.int64),
"aten.argmin": lambda self, dim=None, keepdim=False: self.argmin(dim, keepdim).cast(dtypes.int64),
"aten.asinh": Tensor.asinh,
"aten.mul": Tensor.mul,
"aten.atanh": Tensor.atanh,
@@ -688,6 +669,7 @@ tiny_backend = {**{k:wrap_out(v) for k,v in tiny_backend_out.items()}, **{
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.max.dim": lambda self, dim, keepdim=False: (self.max(dim, keepdim), self.argmax(dim, keepdim).cast(dtype=dtypes.int64)),
"aten.min.dim": lambda self, dim, keepdim=False: (self.min(dim, keepdim), self.argmin(dim, keepdim).cast(dtype=dtypes.int64)),
"aten.cummax": lambda self, dim: ((r := self.cummax(dim))[0], r[1].cast(dtypes.int64)),
"aten.cummin": lambda self, dim: ((r := self.cummin(dim))[0], r[1].cast(dtypes.int64)),
"aten.nonzero": Tensor.nonzero,
extra/torch_backend/test.py
+38
View File
@@ -166,6 +166,15 @@ class TestTorchBackend(unittest.TestCase):
expected = np.array([[1.5, 5.2, 9.0], [13.2, 17.1, 18.4]], dtype=np.float32)
np.testing.assert_equal(y3.cpu().numpy(), expected)
def test_argmax_argmin(self):
a = torch.arange(12, dtype=torch.float32, device=device).reshape(3, 4)
c = a.cpu()
for got, want in [(a.argmax(), c.argmax()), (a.argmin(0), c.argmin(0)), (a.argmax(1, keepdim=True), c.argmax(1, keepdim=True)),
(torch.min(a, 1).indices, torch.min(c, 1).indices), (torch.max(a, 1).indices, torch.max(c, 1).indices),
(torch.min(a, 1).values, torch.min(c, 1).values), (torch.min(a, 1, keepdim=True).indices, torch.min(c, 1, keepdim=True).indices)]:
self.assertEqual(got.dtype, want.dtype) # torch's arg reduces are int64, tinygrad's are int32
np.testing.assert_equal(got.cpu().numpy(), want.numpy())
def test_isfinite(self):
a = torch.ones(4, device=device)
np.testing.assert_equal(torch.isfinite(a).cpu().numpy(), [True, True, True, True])
@@ -373,6 +382,22 @@ class TestTorchBackend(unittest.TestCase):
for bwd_eps in [1e-5, 0.3]:
for got, want in zip(run(device, bwd_eps), run("cpu", bwd_eps)): np.testing.assert_allclose(got, want, atol=1e-4, rtol=1e-3)
def test_groupnorm_backward(self):
def run(dev):
x = torch.arange(24., device=dev).reshape(2, 4, 3).requires_grad_()
w = torch.linspace(0.5, 2.0, 4).to(dev).requires_grad_()
torch.nn.functional.group_norm(x, 2, w, torch.zeros(4, device=dev)).square().sum().backward()
return x.grad.cpu().numpy(), w.grad.cpu().numpy()
for got, want in zip(run(device), run("cpu")): np.testing.assert_allclose(got, want, atol=1e-4, rtol=1e-3)
def test_mse_smooth_l1_loss_backward(self):
def run(dev, loss):
x = torch.arange(4., device=dev).requires_grad_()
loss(x, torch.ones(4, device=dev)).backward()
return x.grad.cpu().numpy()
for loss in [torch.nn.functional.mse_loss, torch.nn.functional.smooth_l1_loss]:
np.testing.assert_allclose(run(device, loss), run("cpu", loss), atol=1e-6)
def test_batchnorm_unsqueeze(self):
bn = torch.nn.BatchNorm2d(4).to(device)
x = torch.randn(8, 4, 3, 3, device=device)
@@ -817,6 +842,19 @@ class TestTorchBackend(unittest.TestCase):
ref = torch.nn.functional.interpolate(a.cpu(), size=(4, 4), mode="bilinear", align_corners=False)
np.testing.assert_allclose(linear.cpu().numpy(), ref.numpy(), rtol=1e-5)
def test_interpolate_bicubic_area(self):
a = torch.arange(32, dtype=torch.float32, device=device).reshape(1, 2, 4, 4)
for mode, scale in [("bicubic", 2.0), ("area", 0.5)]:
ref = torch.nn.functional.interpolate(a.cpu(), scale_factor=scale, mode=mode)
np.testing.assert_allclose(torch.nn.functional.interpolate(a, scale_factor=scale, mode=mode).cpu().numpy(), ref.numpy(), atol=1e-4)
@unittest.expectedFailure
def test_interpolate_bicubic_backward(self):
# the forward comes from a decomposition, but aten::upsample_bicubic2d_backward has none (nor does
# aten::_adaptive_avg_pool2d_backward, for area), so training through these modes needs a real kernel
x = torch.arange(32., dtype=torch.float32, device=device).reshape(1, 2, 4, 4).requires_grad_()
torch.nn.functional.interpolate(x, scale_factor=2.0, mode="bicubic").sum().backward()
@unittest.expectedFailure
def test_interpolate_inexact_scale(self):
# torch forwards the raw scale_factor, Tensor.interpolate recomputes it from output_size, and they disagree here