mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-11-08 13:33:42 +08:00
firestar5683
22a8575199
Keep official openpilot ancestry and dependency licenses. Store resources directly in Git, remove LFS setup requirements, and leave the downloadable big model outside the source tree.
1159 lines
50 KiB
Python
1159 lines
50 KiB
Python
# simple tests
|
|
import unittest
|
|
import torch
|
|
import numpy as np
|
|
from tinygrad.helpers import getenv, GlobalCounters
|
|
if getenv("TINY_BACKEND2"):
|
|
import extra.torch_backend.backend2
|
|
device = "cpu"
|
|
else:
|
|
import extra.torch_backend.backend
|
|
device = "tiny"
|
|
|
|
class TestTorchBackend(unittest.TestCase):
|
|
def test_randperm_generator_out(self):
|
|
n = 10
|
|
out = torch.empty(n, dtype=torch.long, device=device)
|
|
res = torch.randperm(n, out=out).cpu().numpy()
|
|
np.testing.assert_equal(set(res), set(range(n)))
|
|
np.testing.assert_equal(out.cpu().numpy(), res)
|
|
|
|
res2 = torch.randperm(n).cpu().numpy()
|
|
np.testing.assert_equal(set(res2), set(range(n)))
|
|
|
|
def test_numpy_ones(self):
|
|
a = torch.ones(4, device=device)
|
|
np.testing.assert_equal(a.cpu().numpy(), [1,1,1,1])
|
|
|
|
def test_numpy_ones_int32(self):
|
|
a = torch.ones(4, dtype=torch.int32, device=device)
|
|
assert a.dtype == torch.int32
|
|
np.testing.assert_equal(a.cpu().numpy(), [1,1,1,1])
|
|
|
|
def test_plus(self):
|
|
a = torch.ones(4, device=device)
|
|
b = torch.ones(4, device=device)
|
|
c = a+b
|
|
np.testing.assert_equal(c.cpu().numpy(), [2,2,2,2])
|
|
|
|
def test_expand(self):
|
|
a = torch.Tensor([1,2,3,4]).to(device)
|
|
out = a.reshape(4,1).expand(4,4)
|
|
np.testing.assert_equal(out.cpu().numpy(), [[1,1,1,1],[2,2,2,2],[3,3,3,3],[4,4,4,4]])
|
|
|
|
def test_reshape(self):
|
|
a = torch.Tensor([[1,2],[3,4]]).to(device)
|
|
np.testing.assert_equal(a.reshape(4).cpu().numpy(), [1,2,3,4])
|
|
np.testing.assert_equal(a.reshape(2,1,2).cpu().numpy(), [[[1,2]],[[3,4]]])
|
|
np.testing.assert_equal(a.unsqueeze(1).cpu().numpy(), [[[1,2]],[[3,4]]])
|
|
np.testing.assert_equal(a.unsqueeze(1).unsqueeze(1).cpu().numpy(), [[[[1,2]]],[[[3,4]]]])
|
|
np.testing.assert_equal(a.unsqueeze(1).unsqueeze(1).squeeze().cpu().numpy(), [[1,2],[3,4]])
|
|
|
|
def test_permute(self):
|
|
a = torch.Tensor([[1,2],[3,4]]).to(device)
|
|
print(a.stride())
|
|
null = a.permute(0,1)
|
|
perm = a.permute(1,0)
|
|
back = perm.permute(1,0)
|
|
np.testing.assert_equal(a.cpu().numpy(), [[1,2],[3,4]])
|
|
np.testing.assert_equal(null.cpu().numpy(), [[1,2],[3,4]])
|
|
np.testing.assert_equal(perm.cpu().numpy(), [[1,3],[2,4]])
|
|
np.testing.assert_equal(back.cpu().numpy(), [[1,2],[3,4]])
|
|
|
|
def test_shrink(self):
|
|
a = torch.Tensor([1,2,3,4]).to(device)
|
|
np.testing.assert_equal(a[:3].cpu().numpy(), [1,2,3])
|
|
np.testing.assert_equal(a[1:].cpu().numpy(), [2,3,4])
|
|
|
|
def test_as_strided(self):
|
|
a = torch.arange(70, device=device).reshape(1,1,10,7)
|
|
a = a.as_strided((1,1,10,5), (0,0,7,1), storage_offset=0)
|
|
a = a.as_strided((1,1,5,5), (50,50,7,1), storage_offset=21)
|
|
np.testing.assert_equal(a.cpu().numpy().sum(-1), [[[115,150,185,220,255]]])
|
|
|
|
def test_storage_offset_of_computed_tensor(self):
|
|
# a computed result owns its storage, so a slice anywhere in its history must not shift the offset
|
|
a = torch.arange(8., device=device)
|
|
self.assertEqual((a[3:]+1).storage_offset(), 0)
|
|
|
|
def test_storage_offset_through_aliases(self):
|
|
a = torch.arange(8., device=device)[3:]
|
|
self.assertEqual(a.detach().storage_offset(), 3)
|
|
self.assertEqual(a.view(torch.int32).storage_offset(), 3)
|
|
torch.add(torch.ones(5, device=device), torch.ones(5, device=device), out=a)
|
|
self.assertEqual(a.detach().storage_offset(), 3)
|
|
|
|
def test_out_refreshes_views_of_base(self):
|
|
a = torch.zeros(4, device=device)
|
|
v = a[2:]
|
|
torch.add(torch.ones(4, device=device), torch.ones(4, device=device), out=a)
|
|
np.testing.assert_equal(v.cpu().numpy(), [2., 2.])
|
|
|
|
@unittest.expectedFailure # TODO: storage offset assumes a contiguous source, use UOp.contiguous_view_offset
|
|
def test_storage_offset_non_contiguous_source(self):
|
|
a = torch.arange(12., device=device).reshape(3,4)
|
|
self.assertEqual(a.permute(1,0)[1:].storage_offset(), 1)
|
|
self.assertEqual(a.flatten()[3:].flip(0).storage_offset(), 0)
|
|
|
|
def test_as_strided_explicit_zero_offset(self):
|
|
# storage_offset=0 is a real offset, not "unspecified": it must not fall back to the input's own offset
|
|
a = torch.arange(6., device=device)
|
|
np.testing.assert_equal(a[3:].as_strided((2,), (1,), 0).cpu().numpy(), [0,1])
|
|
np.testing.assert_equal(a[3:].as_strided((2,), (1,)).cpu().numpy(), [3,4])
|
|
|
|
def test_empty_strided_default_dtype(self):
|
|
self.assertEqual(torch.empty_strided((2,3), (1,2), device=device).dtype, torch.get_default_dtype())
|
|
|
|
@unittest.expectedFailure # TODO: empty_strided ignores the requested strides, the backend treats everything as contiguous
|
|
def test_empty_strided_honors_strides(self):
|
|
self.assertEqual(tuple(torch.empty_strided((2,3), (1,2), device=device).stride()), (1,2))
|
|
|
|
@unittest.expectedFailure # TODO: torch refuses an out= that overlaps an input, we compute silently
|
|
def test_out_overlapping_input_is_rejected(self):
|
|
x = torch.arange(6., device=device)
|
|
with self.assertRaises(RuntimeError): torch.add(x[:-1], 10, out=x[1:])
|
|
|
|
def test_out_disjoint_input_is_allowed(self):
|
|
# torch permits an out= that shares a base with an input as long as they do not overlap
|
|
x, xc = torch.arange(6., device=device), torch.arange(6.)
|
|
torch.add(x[:3], 10, out=x[3:]); torch.add(xc[:3], 10, out=xc[3:])
|
|
np.testing.assert_equal(x.cpu().numpy(), xc.numpy())
|
|
|
|
def test_plus_inplace(self):
|
|
a = torch.ones(4, device=device)
|
|
b = torch.ones(4, device=device)
|
|
a += b
|
|
a += b
|
|
np.testing.assert_equal(a.cpu().numpy(), [3,3,3,3])
|
|
|
|
def test_exp2(self):
|
|
a = torch.ones(4, device=device)
|
|
b = a.exp2()
|
|
np.testing.assert_equal(b.cpu().numpy(), [2,2,2,2])
|
|
|
|
def test_amax(self):
|
|
x = torch.tensor([[[ 1.5, 2.3, 3.1, 4.7],
|
|
[ 5.2, 6.8, 7.4, 12.9],
|
|
[ 9.0, 12.3, 11.6, 10.1]],
|
|
[[13.2, 16.9, 15.5, 14.1],
|
|
[17.1, 24.9, 19.8, 20.2],
|
|
[21.0, 22.3, 23.6, 18.4]]], device=device)
|
|
|
|
y1 = torch.amax(x)
|
|
expected = np.array([24.9], dtype=np.float32)
|
|
np.testing.assert_equal(y1.cpu().numpy(), expected)
|
|
|
|
y2 = torch.amax(x, dim=(1,2))
|
|
expected = np.array([12.9, 24.9], dtype=np.float32)
|
|
np.testing.assert_equal(y2.cpu().numpy(), expected)
|
|
|
|
y3 = torch.amax(x, dim=2)
|
|
expected = np.array([[4.7, 12.9, 12.3], [16.9, 24.9, 23.6]], dtype=np.float32)
|
|
np.testing.assert_equal(y3.cpu().numpy(), expected)
|
|
|
|
|
|
def test_amin(self):
|
|
x = torch.tensor([[[ 1.5, 2.3, 3.1, 4.7],
|
|
[ 5.2, 6.8, 7.4, 12.9],
|
|
[ 9.0, 12.3, 11.6, 10.1]],
|
|
[[13.2, 16.9, 15.5, 14.1],
|
|
[17.1, 24.9, 19.8, 20.2],
|
|
[21.0, 22.3, 23.6, 18.4]]], device=device)
|
|
|
|
y1 = torch.amin(x)
|
|
expected = np.array([1.5], dtype=np.float32)
|
|
np.testing.assert_equal(y1.cpu().numpy(), expected)
|
|
|
|
y2 = torch.amin(x, dim=(1,2))
|
|
expected = np.array([1.5, 13.2], dtype=np.float32)
|
|
np.testing.assert_equal(y2.cpu().numpy(), expected)
|
|
|
|
y3 = torch.amin(x, dim=2)
|
|
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])
|
|
|
|
def test_eq(self):
|
|
a = torch.ones(4, device=device)
|
|
b = torch.ones(4, device=device)
|
|
c = a == b
|
|
print(c.cpu())
|
|
|
|
def test_maxpool2d_backward(self):
|
|
x = torch.arange(3*3, dtype=torch.float32, device=device).reshape(1, 1, 3, 3).requires_grad_(True)
|
|
torch.nn.functional.max_pool2d(x, kernel_size=2, stride=1).sum().backward()
|
|
np.testing.assert_equal(x.grad.squeeze().cpu().numpy(), [[0, 0, 0], [0, 1, 1], [0, 1, 1]])
|
|
|
|
def test_copy_cast(self):
|
|
x = torch.zeros(4, device=device, dtype=torch.int64)
|
|
y = torch.ones(4, device=device, dtype=torch.float32).to(dtype=torch.int64)
|
|
res1 = x ^ y # an operation that only works on int types
|
|
print(res1.cpu())
|
|
y = y.cpu().float().to(device=device, dtype=torch.int64)
|
|
res2 = x ^ y
|
|
print(res2.cpu())
|
|
|
|
def test_topk(self):
|
|
# test topk return_types
|
|
a = torch.tensor([1, 3, 2, 4], device=device)
|
|
out = torch.topk(a, k=2)
|
|
np.testing.assert_equal(out.values.cpu().numpy(), [4, 3])
|
|
np.testing.assert_equal(out.indices.cpu().numpy(), [3, 1])
|
|
|
|
def test_masked_select(self):
|
|
a = torch.tensor([4, 3, 2, 1], device=device)
|
|
mask = torch.tensor([True, False, True, False], device=device)
|
|
out = torch.masked_select(a, mask)
|
|
np.testing.assert_equal(out.cpu().numpy(), [4, 2])
|
|
mask = torch.tensor(True, device=device)
|
|
out = torch.masked_select(a, mask)
|
|
np.testing.assert_equal(out.cpu().numpy(), [4, 3, 2, 1])
|
|
|
|
def test_isin_tensor_tensor_out(self):
|
|
a = torch.tensor([1, 2, 3], device=device)
|
|
b = torch.tensor([2, 4], device=device)
|
|
expected_base = torch.tensor([False, True, False], device=device)
|
|
for assume_unique in [False, True]:
|
|
for invert, expected in [(False, expected_base), (True, ~expected_base)]:
|
|
out = torch.empty_like(a, dtype=torch.bool)
|
|
res = torch.ops.aten.isin.Tensor_Tensor_out(a, b, invert=invert, assume_unique=assume_unique, out=out)
|
|
np.testing.assert_equal(out.cpu().numpy(), expected.cpu().numpy())
|
|
|
|
def test_uniform(self):
|
|
for torch_dtype in [torch.float32, torch.float16]:
|
|
a = torch.rand(10, 10, device=device, dtype=torch_dtype)
|
|
self.assertEqual(a.dtype, torch_dtype)
|
|
|
|
def test_normal(self):
|
|
for torch_dtype in [torch.float32, torch.float16]:
|
|
a = torch.randn(10, 10, device=device, dtype=torch_dtype)
|
|
self.assertEqual(a.dtype, torch_dtype)
|
|
|
|
def test_equal(self):
|
|
tensor_a = torch.tensor([[1, 2], [3, 4]], device=device)
|
|
tensor_b = torch.tensor([[1, 2], [3, 4]], device=device)
|
|
tensor_c = torch.tensor([[1, 2], [1, 2]], device=device)
|
|
assert torch.equal(tensor_a, tensor_b)
|
|
assert not torch.equal(tensor_a, tensor_c)
|
|
|
|
@unittest.skip("# TODO: this test is slow")
|
|
def test_linalg_svd(self):
|
|
A = torch.randn(5, 5, device=device)
|
|
U, S, Vh = torch.linalg.svd(A)
|
|
np.testing.assert_equal(U.shape, (5,5))
|
|
np.testing.assert_equal(Vh.shape, (5,5))
|
|
np.testing.assert_allclose(torch.dist(A, U @ torch.diag(S) @ Vh).cpu().numpy(), 0, atol=1e-5)
|
|
|
|
A = torch.randn(5, 3, device=device)
|
|
U, S, Vh = torch.linalg.svd(A, full_matrices=False)
|
|
np.testing.assert_equal(U.shape, (5,3))
|
|
np.testing.assert_equal(Vh.shape, (3,3))
|
|
np.testing.assert_allclose(torch.dist(A, U @ torch.diag(S) @ Vh).cpu().numpy(), 0, atol=1e-5)
|
|
|
|
def test_linalg_eigh(self):
|
|
a = torch.tensor([[1, 2], [2, 1]], dtype=torch.float32, device=device)
|
|
w, v = torch.linalg.eigh(a)
|
|
np.testing.assert_equal(w.cpu().numpy(), [-1, 3])
|
|
recon = (v @ torch.diag(w) @ v.T).cpu().numpy()
|
|
np.testing.assert_allclose(recon, a.cpu().numpy(), atol=1e-6)
|
|
|
|
def test_linalg_det(self):
|
|
a = torch.diag(torch.tensor([1,2,3,4,5], dtype = torch.float32, device=device))
|
|
b = torch.linalg.det(a)
|
|
np.testing.assert_equal(b.cpu().numpy(), 120.0)
|
|
|
|
def test_linalg_cross(self):
|
|
a = torch.tensor([[1, 0, 0], [0, 1, 0]], dtype=torch.float32, device=device)
|
|
b = torch.tensor([[0, 0, 1]], dtype=torch.float32, device=device)
|
|
cross = torch.linalg.cross(a, b)
|
|
np.testing.assert_equal(cross.cpu().numpy(), np.array([[0, -1, 0], [1, 0, 0]], dtype=np.float32))
|
|
|
|
def test_scalar_assign(self):
|
|
a = torch.tensor([1, 2, 3], device=device)
|
|
a[1] = 4
|
|
np.testing.assert_equal(a.cpu().numpy(), [1, 4, 3])
|
|
|
|
@unittest.skip("meh")
|
|
def test_str(self):
|
|
a = torch.ones(4, device=device)
|
|
print(str(a))
|
|
|
|
def test_floor_div(self):
|
|
a = torch.tensor([10., 7., 5.], device=device)
|
|
b = torch.tensor([3., 2., 2.], device=device)
|
|
result = a // b
|
|
np.testing.assert_equal(result.cpu().numpy(), [3., 3., 2.])
|
|
|
|
def test_mnist_index(self):
|
|
# from tinygrad.nn.datasets import mnist
|
|
X_train, Y_train = Tensor.randint(60000, 1, 28, 28, dtype='uchar').realize(), Tensor.randint(60000, dtype='uchar').realize()
|
|
GlobalCounters.reset()
|
|
X_train = torch.tensor(X_train.float().numpy(), device=device)
|
|
Y_train = torch.tensor(Y_train.cast('int64').numpy(), device=device)
|
|
samples = torch.randint(0, X_train.shape[0], (32,))
|
|
X,Y = X_train[samples], Y_train[samples]
|
|
X.cpu(), Y.cpu()
|
|
self.assertLessEqual(GlobalCounters.global_ops, 25_000_000)
|
|
|
|
def _test_diagonal(self, *shape):
|
|
a = torch.randn(*shape, dtype=torch.float32, device=device)
|
|
ref = np.diagonal(a.cpu().numpy(), axis1=-2, axis2=-1)
|
|
diag = torch.linalg.diagonal(a)
|
|
np.testing.assert_equal(diag.cpu().numpy(), ref)
|
|
np.testing.assert_equal(diag[-1].cpu().numpy(), ref[-1])
|
|
|
|
def test_diagonal_cube(self): self._test_diagonal(3, 3, 3)
|
|
def test_diagonal_rectangular(self): self._test_diagonal(4, 5, 6)
|
|
def test_diagonal_4d(self): self._test_diagonal(2, 3, 4, 5)
|
|
|
|
def test_pad_circular_simple(self):
|
|
a = torch.arange(4, dtype=torch.float32, device=device).reshape(1,1,2,2)
|
|
padded = torch.nn.functional.pad(a, (1,1,1,1), mode="circular")
|
|
expected = np.array([[[[3.,2.,3.,2.], [1.,0.,1.,0.], [3.,2.,3.,2.], [1.,0.,1.,0.]]]], dtype=np.float32)
|
|
np.testing.assert_allclose(padded.cpu().numpy(), expected)
|
|
|
|
def test_pad_circular_backward(self):
|
|
a = torch.arange(4, dtype=torch.float32, device=device).reshape(1,1,2,2).requires_grad_(True)
|
|
padded = torch.nn.functional.pad(a, (1,1,1,1), mode="circular")
|
|
loss = padded.sum()
|
|
loss.backward()
|
|
expected_grad = np.array([[[[4., 4.], [4., 4.]]]], dtype=np.float32)
|
|
np.testing.assert_allclose(a.grad.cpu().numpy(), expected_grad)
|
|
|
|
|
|
def test_matmul_backward(self):
|
|
x = torch.randn(3, 4, device=device, dtype=torch.float32, requires_grad=True)
|
|
y = torch.randn(4, 5, device=device, dtype=torch.float32, requires_grad=True)
|
|
z = (x @ y).sum()
|
|
z.backward()
|
|
assert x.grad is not None
|
|
assert y.grad is not None
|
|
assert x.grad.shape == x.shape
|
|
assert y.grad.shape == y.shape
|
|
|
|
def test_matmul_broadcast_backward(self):
|
|
x = torch.randn(2, 3, 4, device=device, dtype=torch.float32, requires_grad=True)
|
|
y = torch.randn(4, 5, device=device, dtype=torch.float32, requires_grad=True)
|
|
z = (x @ y).sum()
|
|
z.backward()
|
|
assert x.grad is not None
|
|
assert y.grad is not None
|
|
assert x.grad.shape == x.shape
|
|
assert y.grad.shape == y.shape
|
|
|
|
def test_diag_vector_to_matrix(self):
|
|
vec = torch.tensor([1., 2., 3., 4., 5.], dtype=torch.float32, device=device)
|
|
mat = torch.diag(vec)
|
|
expected = np.diag([1., 2., 3., 4., 5.])
|
|
np.testing.assert_allclose(mat.cpu().numpy(), expected, rtol=1e-5)
|
|
assert mat.shape == (5, 5)
|
|
|
|
def test_diagonal_matrix_to_vector(self):
|
|
mat = torch.tensor([[1., 2., 3.], [4., 5., 6.], [7., 8., 9.]], dtype=torch.float32, device=device)
|
|
vec = torch.linalg.diagonal(mat)
|
|
expected = np.array([1., 5., 9.])
|
|
np.testing.assert_allclose(vec.cpu().numpy(), expected, rtol=1e-5)
|
|
assert vec.shape == (3,)
|
|
|
|
def test_permute_2(self):
|
|
a = torch.randn(2, 3, 4, dtype=torch.float32, device=device)
|
|
b = a.permute(2, 0, 1)
|
|
assert b.shape == (4, 2, 3)
|
|
np.testing.assert_equal(b.cpu().numpy(), a.cpu().numpy().transpose(2, 0, 1))
|
|
|
|
def test_batchnorm_backward_realized_stats(self):
|
|
# the saved stats are a function of input in training, so grad_input must flow through them even when handed in realized.
|
|
# the backward eps is unused in training: torch differentiates the save_invstd it was given
|
|
x0, g0 = torch.randn(8, 4, 3, 3), torch.randn(8, 4, 3, 3)
|
|
def run(dev, bwd_eps):
|
|
x, go = x0.to(dev), g0.to(dev)
|
|
w, b = torch.linspace(0.5, 2.0, 4).to(dev), torch.zeros(4, device=dev)
|
|
rm, rv = torch.zeros(4, device=dev), torch.ones(4, device=dev)
|
|
out, sm, si = torch.ops.aten.native_batch_norm(x, w, b, rm, rv, True, 0.1, 1e-5)
|
|
grads = torch.ops.aten.native_batch_norm_backward(go, x, w, rm, rv, sm.clone().detach(), si.clone().detach(),
|
|
True, bwd_eps, [True,True,True])
|
|
return [t.cpu().numpy() for t in grads]
|
|
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)
|
|
out = bn(x)
|
|
self.assertEqual(out.shape, x.shape)
|
|
|
|
def test_slice_inplace_zero(self):
|
|
a = torch.ones((3, 3), device=device)
|
|
b = a[1:, 1:]
|
|
b.zero_()
|
|
expected = np.array([[1., 1., 1.],
|
|
[1., 0., 0.],
|
|
[1., 0., 0.]])
|
|
np.testing.assert_equal(a.cpu().numpy(), expected)
|
|
|
|
def test_slice_inplace_fill(self):
|
|
a = torch.ones((3, 3), device=device)
|
|
b = a[1:, 1:]
|
|
b.fill_(5.0)
|
|
expected = np.array([[1., 1., 1.],
|
|
[1., 5., 5.],
|
|
[1., 5., 5.]])
|
|
np.testing.assert_equal(a.cpu().numpy(), expected)
|
|
|
|
def test_fill_tensor_value(self):
|
|
a = torch.zeros((2, 2), dtype=torch.float32, device=device)
|
|
value = torch.tensor(3, dtype=torch.int64, device=device)
|
|
a.fill_(value)
|
|
expected = np.full((2, 2), 3, dtype=np.float32)
|
|
np.testing.assert_equal(a.cpu().numpy(), expected)
|
|
|
|
def test_slice_inplace_mul(self):
|
|
a = torch.ones((3, 3), device=device)
|
|
b = a[1:, 1:]
|
|
b *= 2
|
|
expected = np.array([[1., 1., 1.],
|
|
[1., 2., 2.],
|
|
[1., 2., 2.]])
|
|
np.testing.assert_equal(a.cpu().numpy(), expected)
|
|
|
|
def test_permute_slice_zero(self):
|
|
a = torch.ones((3, 3), device=device)
|
|
b = a[1:, 1:].permute(1, 0)
|
|
b.zero_()
|
|
expected = np.array([[1., 1., 1.],
|
|
[1., 0., 0.],
|
|
[1., 0., 0.]])
|
|
np.testing.assert_equal(a.cpu().numpy(), expected)
|
|
|
|
def test_permute_slice_mul(self):
|
|
a = torch.ones((3, 3), device=device)
|
|
b = a[1:, 1:].permute(1, 0)
|
|
b *= 2
|
|
expected = np.array([[1., 1., 1.],
|
|
[1., 2., 2.],
|
|
[1., 2., 2.]])
|
|
np.testing.assert_equal(a.cpu().numpy(), expected)
|
|
|
|
def test_simple_slice_setitem(self):
|
|
a = torch.tensor([10, 20, 30], device=device)
|
|
a[1] = 99
|
|
np.testing.assert_equal(a.cpu().numpy(), [10, 99, 30])
|
|
|
|
def test_2d_slice_setitem(self):
|
|
a = torch.zeros((3, 3), device=device)
|
|
a[1, 2] = 99
|
|
self.assertEqual(a[1, 2].item(), 99)
|
|
self.assertEqual(a.sum().item(), 99)
|
|
|
|
def test_view_copy(self):
|
|
a = torch.tensor([10, 20, 30], device=device)
|
|
view = a[1]
|
|
view.copy_(torch.tensor(88, device=device))
|
|
np.testing.assert_equal(a.cpu().numpy(), [10, 88, 30])
|
|
|
|
def test_diag_2d_input(self):
|
|
a = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]], device=device)
|
|
d = torch.diag(a)
|
|
np.testing.assert_equal(d.cpu().numpy(), [1, 5, 9])
|
|
|
|
def test_diag_1d_input(self):
|
|
a = torch.tensor([1, 2, 3], device=device)
|
|
d = torch.diag(a)
|
|
expected = [[1, 0, 0], [0, 2, 0], [0, 0, 3]]
|
|
np.testing.assert_equal(d.cpu().numpy(), expected)
|
|
|
|
def test_permute_view_tracking(self):
|
|
a = torch.ones((2, 3, 4), device=device)
|
|
b = a.permute(2, 0, 1)
|
|
self.assertEqual(b.shape, (4, 2, 3))
|
|
|
|
def test_detach_view_creation(self):
|
|
a = torch.tensor([1.0, 2.0, 3.0], device=device)
|
|
b = a.detach()
|
|
np.testing.assert_equal(b.cpu().numpy(), [1.0, 2.0, 3.0])
|
|
|
|
def test_view_zero_inplace(self):
|
|
a = torch.ones((4, 4), device=device)
|
|
view = a[1:3, 1:3]
|
|
view.zero_()
|
|
self.assertEqual(view.sum().item(), 0)
|
|
|
|
def test_view_fill_inplace(self):
|
|
a = torch.zeros((4, 4), device=device)
|
|
view = a[1:3, 1:3]
|
|
view.fill_(5)
|
|
self.assertEqual(view.sum().item(), 20)
|
|
|
|
def test_permute_contiguous(self):
|
|
a = torch.tensor([[1, 2], [3, 4]], device=device)
|
|
b = a.permute(1, 0)
|
|
c = b.contiguous()
|
|
expected = [[1, 3], [2, 4]]
|
|
np.testing.assert_equal(c.cpu().numpy(), expected)
|
|
|
|
def test_diag_2d_extract_diagonal(self):
|
|
a = torch.tensor([[1, 2], [3, 4]], device=device)
|
|
result = torch.diag(a)
|
|
np.testing.assert_equal(result.cpu().numpy(), [1, 4])
|
|
|
|
def test_slice_inplace_multiply_offset_preservation(self):
|
|
a = torch.tensor([1, 2, 3], device=device)
|
|
a[1:] *= 2
|
|
np.testing.assert_equal(a.cpu().numpy(), [1, 4, 6])
|
|
|
|
def test_slice_inplace_mul_pattern(self):
|
|
a = torch.tensor([1, 2, 3, 4], device=device)
|
|
a[:2] *= 3
|
|
a[2:] *= 2
|
|
np.testing.assert_equal(a.cpu().numpy(), [3, 6, 6, 8])
|
|
|
|
def test_chained_slice_column(self):
|
|
a = torch.arange(16, dtype=torch.float32, device=device).reshape(4, 4)
|
|
torch_res = a[:, 1:2][:, 0:1].cpu().numpy()
|
|
cpu_res = torch.arange(16, dtype=torch.float32).reshape(4, 4)[:, 1:2][:, 0:1].numpy()
|
|
np.testing.assert_equal(torch_res, cpu_res)
|
|
|
|
def test_slice_with_step(self):
|
|
a = torch.arange(20, dtype=torch.float32, device=device)
|
|
torch_res = a[::2][1:4].cpu().numpy()
|
|
cpu_res = torch.arange(20, dtype=torch.float32)[::2][1:4].numpy()
|
|
np.testing.assert_equal(torch_res, cpu_res)
|
|
|
|
def test_select_out_of_range_dim(self):
|
|
a = torch.arange(12, dtype=torch.int32, device=device).reshape(3, 4)
|
|
with self.assertRaises(IndexError): a.select(5, 0)
|
|
|
|
def test_select_collapses_the_only_dim(self):
|
|
a = torch.arange(3, dtype=torch.int32, device=device)
|
|
self.assertEqual(a.select(0, 1).shape, ())
|
|
np.testing.assert_equal(a.select(0, 1).cpu().numpy(), 1)
|
|
|
|
def test_slice_negative_dim(self):
|
|
a = torch.arange(13, dtype=torch.int32, device=device).repeat(8, 1)
|
|
torch_chunks = a.chunk(3, -1)
|
|
cpu_chunks = torch.arange(13, dtype=torch.int32).repeat(8, 1).chunk(3, -1)
|
|
assert len(torch_chunks) == len(cpu_chunks)
|
|
for i in range(len(torch_chunks)):
|
|
np.testing.assert_equal(torch_chunks[i].cpu().numpy(), cpu_chunks[i].numpy())
|
|
|
|
def test_dot_vector_matrix(self):
|
|
a = torch.arange(65, dtype=torch.float32, device=device)
|
|
b = torch.arange(65*45, dtype=torch.float32, device=device).reshape(65, 45)
|
|
torch_res = a.matmul(b).reshape(-1).cpu().numpy()
|
|
cpu_res = torch.arange(65, dtype=torch.float32).matmul(torch.arange(65*45, dtype=torch.float32).reshape(65, 45)).numpy()
|
|
np.testing.assert_equal(torch_res, cpu_res)
|
|
|
|
def test_alias_passthrough(self):
|
|
a = torch.randn(3, 3, device=device)
|
|
alias_view = torch.ops.aten.alias(a)
|
|
alias_view += 1
|
|
np.testing.assert_equal(a.cpu().numpy(), alias_view.cpu().numpy())
|
|
|
|
def test_split_simple_vector(self):
|
|
a = torch.arange(10, dtype=torch.float32, device=device)
|
|
torch_chunks = a.split([1,4,5])
|
|
cpu_chunks = torch.arange(10, dtype=torch.float32).split([1,4,5])
|
|
for tc, cc in zip(torch_chunks, cpu_chunks):
|
|
np.testing.assert_equal(tc.cpu().numpy(), cc.cpu().numpy())
|
|
|
|
def test_split_matches_torch(self):
|
|
a = torch.arange(10, dtype=torch.float32, device=device)
|
|
torch_chunks = a.split([1,4,5])
|
|
tiny_chunks = [chunk.cpu().numpy() for chunk in torch_chunks]
|
|
cpu_chunks = [torch.arange(10, dtype=torch.float32).split([1,4,5])[i].numpy() for i in range(3)]
|
|
for tr, cr in zip(tiny_chunks, cpu_chunks): np.testing.assert_equal(tr, cr)
|
|
|
|
def test_sum_matches_torch(self):
|
|
a = torch.arange(6, dtype=torch.float32, device=device).reshape(2,3)
|
|
torch_res = a.sum().cpu().numpy()
|
|
cpu_res = torch.arange(6, dtype=torch.float32).reshape(2,3).sum().numpy()
|
|
np.testing.assert_equal(torch_res, cpu_res)
|
|
|
|
def test_count_nonzero(self):
|
|
a = torch.tensor([[0, 2, 0], [3, 0, -4]], device=device)
|
|
torch_res = a.count_nonzero().cpu().numpy()
|
|
cpu_res = torch.tensor([[0, 2, 0], [3, 0, -4]]).count_nonzero().numpy()
|
|
np.testing.assert_equal(torch_res, cpu_res)
|
|
|
|
def test_view_matches_torch(self):
|
|
a = torch.arange(6, dtype=torch.float32, device=device)
|
|
torch_res = a.view(2, 3).cpu().numpy()
|
|
cpu_res = torch.arange(6, dtype=torch.float32).view(2, 3).numpy()
|
|
np.testing.assert_equal(torch_res, cpu_res)
|
|
|
|
def test_view_zero_with_indices(self):
|
|
a = torch.tensor([1, 2, 3, 4], device=device)
|
|
a[1:3].zero_()
|
|
np.testing.assert_equal(a.cpu().numpy(), [1, 0, 0, 4])
|
|
|
|
def test_view_fill_with_indices(self):
|
|
a = torch.tensor([1, 2, 3, 4], device=device)
|
|
a[::2].fill_(9)
|
|
np.testing.assert_equal(a.cpu().numpy(), [9, 2, 9, 4])
|
|
|
|
def test_nested_slice_inplace_ops(self):
|
|
a = torch.tensor([1, 2, 3, 4, 5, 6], device=device)
|
|
a[:3] += 10
|
|
a[3:] *= 2
|
|
np.testing.assert_equal(a.cpu().numpy(), [11, 12, 13, 8, 10, 12])
|
|
|
|
def test_diag_1d(self):
|
|
a = torch.tensor([1, 2, 3], device=device)
|
|
result = torch.diag(a)
|
|
expected = [[1, 0, 0], [0, 2, 0], [0, 0, 3]]
|
|
np.testing.assert_equal(result.cpu().numpy(), expected)
|
|
|
|
def test_diag_backward(self):
|
|
a = torch.randn(5, dtype=torch.float32, device=device, requires_grad=True)
|
|
b = torch.diag(a)
|
|
b.sum().backward()
|
|
assert a.grad is not None
|
|
|
|
def test_diagonal(self):
|
|
a = torch.tensor([[1., 2., 3.], [4., 5., 6.], [7., 8., 9.]], dtype=torch.float32, device=device, requires_grad=True)
|
|
b = torch.diagonal(a)
|
|
expected = torch.tensor([1., 5., 9.], dtype=torch.float32)
|
|
self.assertEqual(b.shape, (3,))
|
|
np.testing.assert_allclose(b.detach().cpu().numpy(), expected.numpy(), rtol=1e-5)
|
|
|
|
def test_diagonal_backward(self):
|
|
a = torch.randn(5, 5, dtype=torch.float32, device=device, requires_grad=True)
|
|
b = torch.diagonal(a)
|
|
b.sum().backward()
|
|
assert a.grad is not None
|
|
|
|
def test_expand_backward(self):
|
|
a = torch.randn(4, 3, 1, 6, dtype=torch.float32, device=device, requires_grad=True)
|
|
b = a.expand(4, 3, 2, 6)
|
|
b.sum().backward()
|
|
assert a.grad is not None
|
|
|
|
def test_einsum_backward(self):
|
|
a = torch.randn(10, 10, dtype=torch.float32, device=device, requires_grad=True)
|
|
b = torch.einsum('ij->ji', a)
|
|
b.sum().backward()
|
|
assert a.grad is not None
|
|
|
|
def test_diag_backward_gradient_values(self):
|
|
a = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32, device=device, requires_grad=True)
|
|
b = torch.diag(a)
|
|
loss = b.sum()
|
|
loss.backward()
|
|
expected_grad = torch.ones(3, dtype=torch.float32)
|
|
np.testing.assert_allclose(a.grad.cpu().numpy(), expected_grad.numpy(), rtol=1e-5)
|
|
|
|
def test_diag_backward_gradient_values_2d_to_1d(self):
|
|
a = torch.tensor([[1.0, 2.0, 3.0],
|
|
[4.0, 5.0, 6.0],
|
|
[7.0, 8.0, 9.0]], dtype=torch.float32, device=device, requires_grad=True)
|
|
b = torch.diagonal(a)
|
|
loss = b.sum()
|
|
loss.backward()
|
|
expected_grad = torch.tensor([[1.0, 0.0, 0.0],
|
|
[0.0, 1.0, 0.0],
|
|
[0.0, 0.0, 1.0]], dtype=torch.float32)
|
|
np.testing.assert_allclose(a.grad.cpu().numpy(), expected_grad.numpy(), rtol=1e-5)
|
|
|
|
def test_expand_backward_gradient_values(self):
|
|
a = torch.tensor([[1.0], [2.0], [3.0]], dtype=torch.float32, device=device, requires_grad=True)
|
|
b = a.expand(3, 4)
|
|
loss = b.sum()
|
|
loss.backward()
|
|
expected_grad = torch.tensor([[4.0], [4.0], [4.0]], dtype=torch.float32)
|
|
np.testing.assert_allclose(a.grad.cpu().numpy(), expected_grad.numpy(), rtol=1e-5)
|
|
|
|
def test_expand_backward_with_leading_dims(self):
|
|
a = torch.tensor([[1.0, 2.0]], dtype=torch.float32, device=device, requires_grad=True)
|
|
b = a.expand(3, 1, 2)
|
|
loss = b.sum()
|
|
loss.backward()
|
|
expected_grad = torch.tensor([[3.0, 3.0]], dtype=torch.float32)
|
|
np.testing.assert_allclose(a.grad.cpu().numpy(), expected_grad.numpy(), rtol=1e-5)
|
|
|
|
def test_diag_2d_to_1d_backward(self):
|
|
a = torch.tensor([[1.0, 2.0], [3.0, 4.0]], dtype=torch.float32, device=device, requires_grad=True)
|
|
b = torch.diag(a)
|
|
loss = b.sum()
|
|
loss.backward()
|
|
expected_grad = torch.tensor([[1.0, 0.0], [0.0, 1.0]], dtype=torch.float32)
|
|
np.testing.assert_allclose(a.grad.cpu().numpy(), expected_grad.numpy(), rtol=1e-5)
|
|
|
|
def test_expand_complex_backward(self):
|
|
a = torch.tensor([[[1.0, 2.0]]], dtype=torch.float32, device=device, requires_grad=True)
|
|
b = a.expand(2, 3, 2)
|
|
loss = b.sum()
|
|
loss.backward()
|
|
expected_grad = torch.tensor([[[6.0, 6.0]]], dtype=torch.float32)
|
|
np.testing.assert_allclose(a.grad.cpu().numpy(), expected_grad.numpy(), rtol=1e-5)
|
|
|
|
def test_diag_backward_with_scaling(self):
|
|
a = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32, device=device, requires_grad=True)
|
|
b = torch.diag(a)
|
|
loss = (b * torch.tensor([[2.0, 0.0, 0.0],
|
|
[0.0, 3.0, 0.0],
|
|
[0.0, 0.0, 4.0]], device=device)).sum()
|
|
loss.backward()
|
|
expected_grad = torch.tensor([2.0, 3.0, 4.0], dtype=torch.float32)
|
|
np.testing.assert_allclose(a.grad.cpu().numpy(), expected_grad.numpy(), rtol=1e-5)
|
|
|
|
def test_repeat_basic(self):
|
|
a = torch.tensor([1, 2, 3], dtype=torch.float32, device=device)
|
|
b = a.repeat(2, 1)
|
|
expected = torch.tensor([[1, 2, 3], [1, 2, 3]], dtype=torch.float32)
|
|
np.testing.assert_equal(b.cpu().numpy(), expected.numpy())
|
|
|
|
def test_repeat_multidim(self):
|
|
a = torch.arange(6, dtype=torch.float32, device=device).reshape(2, 3)
|
|
b = a.repeat(2, 3)
|
|
expected = torch.arange(6, dtype=torch.float32).reshape(2, 3).repeat(2, 3)
|
|
np.testing.assert_equal(b.cpu().numpy(), expected.numpy())
|
|
|
|
def test_repeat_backward(self):
|
|
a = torch.tensor([[1.0, 2.0]], dtype=torch.float32, device=device, requires_grad=True)
|
|
b = a.repeat(3, 2)
|
|
loss = b.sum()
|
|
loss.backward()
|
|
expected_grad = torch.tensor([[6.0, 6.0]], dtype=torch.float32)
|
|
np.testing.assert_allclose(a.grad.cpu().numpy(), expected_grad.numpy(), rtol=1e-5)
|
|
|
|
def test_cumsum_1d(self):
|
|
a = torch.tensor([1, 2, 3, 4], dtype=torch.float32, device=device)
|
|
b = torch.cumsum(a, dim=0)
|
|
expected = torch.tensor([1, 3, 6, 10], dtype=torch.float32)
|
|
np.testing.assert_equal(b.cpu().numpy(), expected.numpy())
|
|
|
|
def test_cumsum_2d(self):
|
|
a = torch.arange(12, dtype=torch.float32, device=device).reshape(3, 4)
|
|
b = torch.cumsum(a, dim=0)
|
|
expected = torch.arange(12, dtype=torch.float32).reshape(3, 4).cumsum(dim=0)
|
|
np.testing.assert_equal(b.cpu().numpy(), expected.numpy())
|
|
|
|
c = torch.cumsum(a, dim=1)
|
|
expected = torch.arange(12, dtype=torch.float32).reshape(3, 4).cumsum(dim=1)
|
|
np.testing.assert_equal(c.cpu().numpy(), expected.numpy())
|
|
|
|
def test_cumsum_backward(self):
|
|
a = torch.tensor([1.0, 2.0, 3.0, 4.0], dtype=torch.float32, device=device, requires_grad=True)
|
|
b = torch.cumsum(a, dim=0)
|
|
loss = b.sum()
|
|
loss.backward()
|
|
expected_grad = torch.tensor([4.0, 3.0, 2.0, 1.0], dtype=torch.float32)
|
|
np.testing.assert_allclose(a.grad.cpu().numpy(), expected_grad.numpy(), rtol=1e-5)
|
|
|
|
def test_constant_pad_nd_1d(self):
|
|
a = torch.tensor([1, 2, 3], dtype=torch.float32, device=device)
|
|
b = torch.nn.functional.pad(a, (1, 2), mode='constant', value=0)
|
|
expected = torch.tensor([0, 1, 2, 3, 0, 0], dtype=torch.float32)
|
|
np.testing.assert_equal(b.cpu().numpy(), expected.numpy())
|
|
|
|
def test_constant_pad_nd_2d(self):
|
|
a = torch.arange(6, dtype=torch.float32, device=device).reshape(2, 3)
|
|
b = torch.nn.functional.pad(a, (1, 1, 1, 1), mode='constant', value=0)
|
|
expected = torch.nn.functional.pad(torch.arange(6, dtype=torch.float32).reshape(2, 3), (1, 1, 1, 1), mode='constant', value=0)
|
|
np.testing.assert_equal(b.cpu().numpy(), expected.numpy())
|
|
|
|
def test_constant_pad_nd_2d_backward(self):
|
|
a = torch.tensor([[1.0, 2.0], [3.0, 4.0]], dtype=torch.float32, device=device, requires_grad=True)
|
|
b = torch.nn.functional.pad(a, (1, 1, 1, 1), mode='constant', value=0)
|
|
loss = b.sum()
|
|
loss.backward()
|
|
expected_grad = torch.ones((2, 2), dtype=torch.float32)
|
|
np.testing.assert_allclose(a.grad.cpu().numpy(), expected_grad.numpy(), rtol=1e-5)
|
|
|
|
def test_negative_strides_cumsum_backward(self):
|
|
a = torch.randn(5, device=device, requires_grad=True)
|
|
b = torch.cumsum(a, dim=0)
|
|
b.sum().backward()
|
|
grad = a.grad.cpu().numpy()
|
|
self.assertEqual(len(grad), 5)
|
|
|
|
def test_cumsum_fix_gradient_values(self):
|
|
a = torch.tensor([1.0, 2.0, 3.0, 4.0], dtype=torch.float32, device=device, requires_grad=True)
|
|
b = torch.cumsum(a, dim=0)
|
|
loss = b.sum()
|
|
loss.backward()
|
|
expected = np.array([4.0, 3.0, 2.0, 1.0])
|
|
np.testing.assert_allclose(a.grad.cpu().numpy(), expected, rtol=1e-5)
|
|
|
|
def test_cumsum_arange_large(self):
|
|
# Tests cumsum with an unrealized arange input with size > 512 (the split threshold)
|
|
# This exercises the _split_cumalu path which uses a two-stage algorithm
|
|
for size in [513, 1022]:
|
|
a = torch.arange(size, dtype=torch.float32, device=device)
|
|
result = torch.cumsum(a, dim=0)
|
|
expected = torch.arange(size, dtype=torch.float32).cumsum(dim=0)
|
|
np.testing.assert_allclose(result.cpu().numpy(), expected.numpy(), rtol=1e-5)
|
|
|
|
def test_diag_1d_to_2d(self):
|
|
a = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32, device=device, requires_grad=True)
|
|
b = torch.diag(a)
|
|
expected = [[1, 0, 0], [0, 2, 0], [0, 0, 3]]
|
|
np.testing.assert_equal(b.detach().cpu().numpy(), expected)
|
|
|
|
def test_diag_2d_to_1d(self):
|
|
c = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]], dtype=torch.float32, device=device)
|
|
d = torch.diag(c)
|
|
np.testing.assert_equal(d.cpu().numpy(), [1, 5, 9])
|
|
|
|
def test_biased_conv2d(self):
|
|
# Test case for two sequential conv2d with same weights/bias and ReLU in between, this is as special case from test_ops.py
|
|
torch.manual_seed(0)
|
|
C = 8
|
|
x_cpu = torch.randn(1, C, 5, 5, requires_grad=True)
|
|
w_cpu = torch.randn(C, C, 1, 1, requires_grad=True)
|
|
b_cpu = torch.randn(C, requires_grad=True)
|
|
x_tiny = x_cpu.detach().to(device).requires_grad_(True)
|
|
w_tiny = w_cpu.detach().to(device).requires_grad_(True)
|
|
b_tiny = b_cpu.detach().to(device).requires_grad_(True)
|
|
out_cpu = torch.nn.functional.conv2d(torch.nn.functional.conv2d(x_cpu, w_cpu, b_cpu).relu(), w_cpu, b_cpu)
|
|
out_tiny = torch.nn.functional.conv2d(torch.nn.functional.conv2d(x_tiny, w_tiny, b_tiny).relu(), w_tiny, b_tiny)
|
|
grad_out = torch.randn_like(out_cpu)
|
|
out_cpu.backward(grad_out)
|
|
out_tiny.backward(grad_out.to(device))
|
|
np.testing.assert_allclose(x_tiny.grad.cpu().numpy(), x_cpu.grad.numpy(), atol=1e-4, rtol=1e-3)
|
|
np.testing.assert_allclose(w_tiny.grad.cpu().numpy(), w_cpu.grad.numpy(), atol=1e-4, rtol=1e-3)
|
|
np.testing.assert_allclose(b_tiny.grad.cpu().numpy(), b_cpu.grad.numpy(), atol=1e-4, rtol=1e-3)
|
|
|
|
def test_write_through_detach_of_unrealized(self):
|
|
a = torch.empty(4, device=device)
|
|
a.detach().fill_(3)
|
|
np.testing.assert_equal(a.cpu().numpy(), [3, 3, 3, 3])
|
|
|
|
def test_square_transpose_inplace(self):
|
|
# a same-shape transpose is not a reshape: writing the transposed values straight back would scramble the base
|
|
a = torch.tensor([[0., 1., 2.], [3., 4., 5.], [6., 7., 8.]], device=device)
|
|
a.transpose(0, 1).add_(100)
|
|
np.testing.assert_equal(a.cpu().numpy(), [[100., 101., 102.], [103., 104., 105.], [106., 107., 108.]])
|
|
|
|
def test_interpolate(self):
|
|
a = torch.arange(4, dtype=torch.float32, device=device).reshape(1, 1, 2, 2)
|
|
nearest = torch.nn.functional.interpolate(a, scale_factor=2.0)
|
|
np.testing.assert_equal(nearest.cpu().numpy()[0, 0], [[0, 0, 1, 1], [0, 0, 1, 1], [2, 2, 3, 3], [2, 2, 3, 3]])
|
|
linear = torch.nn.functional.interpolate(a, size=(4, 4), mode="bilinear", align_corners=False)
|
|
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
|
|
a = torch.arange(6, dtype=torch.float32, device=device).reshape(1, 1, 2, 3)
|
|
tiny = torch.nn.functional.interpolate(a, scale_factor=2.5, mode="bilinear")
|
|
ref = torch.nn.functional.interpolate(a.cpu(), scale_factor=2.5, mode="bilinear")
|
|
np.testing.assert_allclose(tiny.cpu().numpy(), ref.numpy(), rtol=1e-5)
|
|
|
|
def test_logical_or_xor(self):
|
|
a = torch.tensor([True, True, False, False], device=device)
|
|
b = torch.tensor([True, False, True, False], device=device)
|
|
np.testing.assert_equal(torch.logical_or(a, b).cpu().numpy(), [True, True, True, False])
|
|
np.testing.assert_equal(torch.logical_xor(a, b).cpu().numpy(), [False, True, True, False])
|
|
# bool-valued whatever the input dtype, so this is not | and ^
|
|
i, j = torch.tensor([2, 0, 5, 0], device=device), torch.tensor([0, 0, 1, 1], device=device)
|
|
np.testing.assert_equal(torch.logical_or(i, j).cpu().numpy(), [True, False, True, True])
|
|
np.testing.assert_equal(torch.logical_xor(i, j).cpu().numpy(), [True, False, False, True])
|
|
|
|
def test_slice_scatter(self):
|
|
# the scatters are functional: they return a new tensor and must leave the one they were given alone
|
|
a = torch.arange(12, dtype=torch.float32, device=device).reshape(3, 4)
|
|
out = torch.slice_scatter(a, torch.ones(1, 4, device=device), 0, 0, 1)
|
|
np.testing.assert_equal(out.cpu().numpy(), [[1, 1, 1, 1], [4, 5, 6, 7], [8, 9, 10, 11]])
|
|
np.testing.assert_equal(a.cpu().numpy(), np.arange(12, dtype=np.float32).reshape(3, 4))
|
|
|
|
def test_slice_scatter_casts_src(self):
|
|
a = torch.zeros(3, 4, device=device)
|
|
out = torch.slice_scatter(a, torch.ones(1, 4, dtype=torch.int32, device=device), 0, 0, 1)
|
|
self.assertEqual(out.dtype, torch.float32)
|
|
np.testing.assert_equal(out.cpu().numpy()[0], np.ones(4, dtype=np.float32))
|
|
|
|
def test_select_scatter(self):
|
|
a = torch.arange(12, dtype=torch.float32, device=device).reshape(3, 4)
|
|
out = torch.select_scatter(a, torch.ones(4, device=device), 0, 1)
|
|
np.testing.assert_equal(out.cpu().numpy(), [[0, 1, 2, 3], [1, 1, 1, 1], [8, 9, 10, 11]])
|
|
|
|
def test_diagonal_scatter(self):
|
|
a = torch.zeros(3, 3, device=device)
|
|
out = torch.diagonal_scatter(a, torch.arange(3, dtype=torch.float32, device=device))
|
|
np.testing.assert_equal(out.cpu().numpy(), np.diag([0., 1., 2.]))
|
|
np.testing.assert_equal(a.cpu().numpy(), np.zeros((3, 3), dtype=np.float32))
|
|
|
|
def test_copy_functional(self):
|
|
# without an impl this segfaults rather than fails: a regression here takes the whole run down
|
|
a = torch.arange(4, dtype=torch.float32, device=device)
|
|
out = torch.ops.aten.copy(a, torch.zeros(4, device=device))
|
|
np.testing.assert_equal(out.cpu().numpy(), [0., 0., 0., 0.])
|
|
np.testing.assert_equal(a.cpu().numpy(), [0., 1., 2., 3.])
|
|
|
|
from tinygrad import Tensor
|
|
class TestBackendHelpers(unittest.TestCase):
|
|
def test_unwrap_rejects_foreign_tensor(self):
|
|
# unwrap casts to the tiny impl, so a tensor from another backend must be refused rather than reinterpreted
|
|
with self.assertRaises(RuntimeError): extra.torch_backend.backend.unwrap(torch.ones(4))
|
|
|
|
def test_update_metadata_rejects_foreign_tensor(self):
|
|
# resizing a tensor we don't own would expose memory past its allocation
|
|
t = torch.ones(4)
|
|
with self.assertRaises(RuntimeError): extra.torch_backend.backend.mod.update_metadata(t, [8], [1], 0)
|
|
self.assertEqual(t.shape, (4,))
|
|
|
|
def test_unwrap_parameter_and_detached(self):
|
|
# nn.Parameter and detach rebuild the base OpaqueTensorImpl, which unwrap still has to accept
|
|
extra.torch_backend.backend.unwrap(torch.nn.Parameter(torch.ones(4, device="tiny")))
|
|
extra.torch_backend.backend.unwrap(torch.ones(4, device="tiny").detach())
|
|
|
|
def test_calculate_storage_offset_no_shrink(self):
|
|
t = Tensor.ones(3, 4)
|
|
assert extra.torch_backend.backend.calculate_storage_offset(t) == 0
|
|
|
|
def test_calculate_storage_offset_with_shrink(self):
|
|
t = Tensor.ones(10, 10)[2:5, 3:7]
|
|
# strides for (10, 10) are [10, 1]
|
|
# offset = 2*10 + 3*1 = 23
|
|
assert extra.torch_backend.backend.calculate_storage_offset(t) == 23
|
|
|
|
def test_calculate_storage_offset_multiple_shrinks(self):
|
|
t = Tensor.ones(5, 6, 7)[1:3, 2:4, 3:5]
|
|
# strides for (5, 6, 7) are [42, 7, 1]
|
|
# offset = 1*42 + 2*7 + 3*1 = 42 + 14 + 3 = 59
|
|
assert extra.torch_backend.backend.calculate_storage_offset(t) == 59
|
|
|
|
def test_calculate_storage_offset_with_reshape(self):
|
|
t = Tensor.ones(10, 10)
|
|
orig_offset = extra.torch_backend.backend.calculate_storage_offset(t)
|
|
assert orig_offset == 0
|
|
t = t.reshape(100)
|
|
assert extra.torch_backend.backend.calculate_storage_offset(t) == orig_offset
|
|
|
|
def test_slice_values_match_torch(self):
|
|
torch_cpu = torch.arange(100, dtype=torch.float32).reshape(10, 10)
|
|
torch_tiny = torch_cpu.to(device)
|
|
sliced_cpu = torch_cpu[2:5, 3:7]
|
|
sliced_tiny = torch_tiny[2:5, 3:7]
|
|
np.testing.assert_equal(sliced_tiny.cpu().numpy(), sliced_cpu.numpy())
|
|
|
|
def test_slice_values_match_torch_3d(self):
|
|
torch_cpu_3d = torch.arange(210, dtype=torch.float32).reshape(5, 6, 7)
|
|
torch_tiny_3d = torch_cpu_3d.to(device)
|
|
sliced_cpu_3d = torch_cpu_3d[1:3, 2:4, 3:5]
|
|
sliced_tiny_3d = torch_tiny_3d[1:3, 2:4, 3:5]
|
|
np.testing.assert_equal(sliced_tiny_3d.cpu().numpy(), sliced_cpu_3d.numpy())
|
|
|
|
def test_topk_out(self):
|
|
a = torch.tensor([1, 3, 2, 4], device=device)
|
|
values = torch.empty(2, device=device)
|
|
indices = torch.empty(2, dtype=torch.int64, device=device)
|
|
ret_values, ret_indices = torch.topk(a, k=2, out=(values, indices))
|
|
np.testing.assert_equal(values.cpu().numpy(), [4, 3])
|
|
np.testing.assert_equal(indices.cpu().numpy(), [3, 1])
|
|
assert ret_values is values
|
|
assert ret_indices is indices
|
|
|
|
def test_sort_out(self):
|
|
a = torch.tensor([3, 1, 4, 2], device=device)
|
|
values = torch.empty(4, device=device)
|
|
indices = torch.empty(4, dtype=torch.int64, device=device)
|
|
ret_values, ret_indices = torch.sort(a, out=(values, indices))
|
|
np.testing.assert_equal(values.cpu().numpy(), [1, 2, 3, 4])
|
|
np.testing.assert_equal(indices.cpu().numpy(), [1, 3, 0, 2])
|
|
assert ret_values is values
|
|
assert ret_indices is indices
|
|
|
|
def test_cat_out(self):
|
|
a = torch.tensor([1, 2], device=device)
|
|
b = torch.tensor([3, 4], device=device)
|
|
out = torch.empty(4, device=device)
|
|
ret = torch.cat([a, b], out=out)
|
|
np.testing.assert_equal(out.cpu().numpy(), [1, 2, 3, 4])
|
|
assert ret is out
|
|
|
|
def test_cat_out_empty_1d(self):
|
|
# Test tiny and cpu to show test passes on torch cpu
|
|
for test_device in device, "cpu":
|
|
a = torch.tensor([], device=device)
|
|
b = torch.tensor([1, 2, 3, 4], device=device).reshape((2, 2))
|
|
out = torch.empty((2, 2), device=device)
|
|
for dim in 0, 1, -1, -2:
|
|
ret = torch.cat([a, b], out=out, dim=dim)
|
|
np.testing.assert_equal(out.cpu().numpy(), [[1, 2], [3, 4]])
|
|
assert ret is out
|
|
|
|
def test_cat_all_empty(self):
|
|
for test_device in device, "cpu":
|
|
a = torch.tensor([], device=device)
|
|
out = torch.empty((0,), device=device)
|
|
for dim in 0, -1:
|
|
ret = torch.cat([a, a], out=out, dim=dim)
|
|
np.testing.assert_equal(out.cpu().numpy(), [])
|
|
assert ret is out
|
|
|
|
def test_scatter_add_out(self):
|
|
src = torch.tensor([[1, 2, 3], [4, 5, 6]], device=device, dtype=torch.float32)
|
|
index = torch.tensor([[0, 1, 2], [0, 1, 2]], device=device)
|
|
input = torch.zeros(3, 3, device=device, dtype=torch.float32)
|
|
out = torch.zeros(3, 3, device=device, dtype=torch.float32)
|
|
ret = torch.scatter_add(input, 0, index, src, out=out)
|
|
expected = torch.tensor([[5, 0, 0], [0, 7, 0], [0, 0, 9]], dtype=torch.float32)
|
|
np.testing.assert_allclose(out.cpu().numpy(), expected.cpu().numpy())
|
|
assert ret is out
|
|
|
|
def test_floor_divide_inplace_identity(self):
|
|
x = torch.tensor([10, 20, 30, 40], dtype=torch.int32, device=device)
|
|
y = torch.tensor([2, 4, 5, 8], dtype=torch.int32, device=device)
|
|
ret = x.floor_divide_(y)
|
|
assert ret is x
|
|
np.testing.assert_equal(x.cpu().numpy(), [5, 5, 6, 5])
|
|
|
|
def test_lshift_inplace_identity(self):
|
|
x = torch.tensor([1, 2, 3, 4], dtype=torch.int32, device=device)
|
|
ret = x.__ilshift__(2)
|
|
assert ret is x
|
|
np.testing.assert_equal(x.cpu().numpy(), [4, 8, 12, 16])
|
|
|
|
def test_rshift_inplace_identity(self):
|
|
x = torch.tensor([16, 32, 48, 64], dtype=torch.int32, device=device)
|
|
ret = x.__irshift__(2)
|
|
assert ret is x
|
|
np.testing.assert_equal(x.cpu().numpy(), [4, 8, 12, 16])
|
|
|
|
def test_relu_inplace_identity(self):
|
|
x = torch.tensor([-1.0, 2.0, -3.0, 4.0], device=device)
|
|
ret = x.relu_()
|
|
assert ret is x
|
|
np.testing.assert_equal(x.cpu().numpy(), [0.0, 2.0, 0.0, 4.0])
|
|
|
|
def test_random_inplace_identity(self):
|
|
x = torch.zeros(10, dtype=torch.int32, device=device)
|
|
ret = x.random_()
|
|
assert ret is x
|
|
assert x.shape == (10,)
|
|
|
|
def test_random_from_inplace_identity(self):
|
|
x = torch.zeros(10, dtype=torch.int32, device=device)
|
|
ret = x.random_(5, 10)
|
|
assert ret is x
|
|
# values should be in range [5, 10)
|
|
assert torch.all(x >= 5).item() and torch.all(x < 10).item()
|
|
|
|
def test_uniform_inplace_identity(self):
|
|
x = torch.zeros(10, device=device)
|
|
ret = x.uniform_(0.0, 1.0)
|
|
assert ret is x
|
|
# values should be in range [0, 1)
|
|
assert torch.all(x >= 0.0).item() and torch.all(x < 1.0).item()
|
|
|
|
def test_normal_inplace_identity(self):
|
|
x = torch.zeros(100, device=device)
|
|
ret = x.normal_(0.0, 1.0)
|
|
assert ret is x
|
|
# just check that values changed from zeros
|
|
assert not torch.all(x == 0.0).item()
|
|
|
|
def test_logical_or_inplace_identity(self):
|
|
x = torch.tensor([True, False, True, False], device=device)
|
|
y = torch.tensor([False, False, True, True], device=device)
|
|
ret = x.logical_or_(y)
|
|
assert ret is x
|
|
np.testing.assert_equal(x.cpu().numpy(), [True, False, True, True])
|
|
|
|
def test_masked_fill_scalar_inplace_identity(self):
|
|
x = torch.tensor([1.0, 2.0, 3.0, 4.0], device=device)
|
|
mask = torch.tensor([True, False, True, False], device=device)
|
|
ret = x.masked_fill_(mask, 0.0)
|
|
assert ret is x
|
|
np.testing.assert_equal(x.cpu().numpy(), [0.0, 2.0, 0.0, 4.0])
|
|
|
|
def test_masked_fill_tensor_inplace_identity(self):
|
|
x = torch.tensor([1.0, 2.0, 3.0, 4.0], device=device)
|
|
mask = torch.tensor([True, False, True, False], device=device)
|
|
value = torch.tensor(99.0, device=device)
|
|
ret = x.masked_fill_(mask, value)
|
|
assert ret is x
|
|
np.testing.assert_equal(x.cpu().numpy(), [99.0, 2.0, 99.0, 4.0])
|
|
|
|
def test_zero_inplace_identity(self):
|
|
x = torch.tensor([1.0, 2.0, 3.0, 4.0], device=device)
|
|
ret = x.zero_()
|
|
assert ret is x
|
|
np.testing.assert_equal(x.cpu().numpy(), [0.0, 0.0, 0.0, 0.0])
|
|
|
|
def test_fill_scalar_inplace_identity(self):
|
|
x = torch.tensor([1.0, 2.0, 3.0, 4.0], device=device)
|
|
ret = x.fill_(5.0)
|
|
assert ret is x
|
|
np.testing.assert_equal(x.cpu().numpy(), [5.0, 5.0, 5.0, 5.0])
|
|
|
|
def test_fill_tensor_inplace_identity(self):
|
|
x = torch.tensor([1.0, 2.0, 3.0, 4.0], device=device)
|
|
value = torch.tensor(7.0, device=device)
|
|
ret = x.fill_(value)
|
|
assert ret is x
|
|
np.testing.assert_equal(x.cpu().numpy(), [7.0, 7.0, 7.0, 7.0])
|
|
|
|
def test_add_tensor_inplace_identity(self):
|
|
x = torch.tensor([1.0, 2.0, 3.0, 4.0], device=device)
|
|
y = torch.tensor([10.0, 20.0, 30.0, 40.0], device=device)
|
|
ret = x.add_(y)
|
|
assert ret is x
|
|
np.testing.assert_equal(x.cpu().numpy(), [11.0, 22.0, 33.0, 44.0])
|
|
|
|
def test_add_scalar_inplace_identity(self):
|
|
x = torch.tensor([1.0, 2.0, 3.0, 4.0], device=device)
|
|
ret = x.add_(10.0)
|
|
assert ret is x
|
|
np.testing.assert_equal(x.cpu().numpy(), [11.0, 12.0, 13.0, 14.0])
|
|
|
|
def test_mul_tensor_inplace_identity(self):
|
|
x = torch.tensor([1.0, 2.0, 3.0, 4.0], device=device)
|
|
y = torch.tensor([2.0, 3.0, 4.0, 5.0], device=device)
|
|
ret = x.mul_(y)
|
|
assert ret is x
|
|
np.testing.assert_equal(x.cpu().numpy(), [2.0, 6.0, 12.0, 20.0])
|
|
|
|
def test_mul_scalar_inplace_identity(self):
|
|
x = torch.tensor([1.0, 2.0, 3.0, 4.0], device=device)
|
|
ret = x.mul_(2.0)
|
|
assert ret is x
|
|
np.testing.assert_equal(x.cpu().numpy(), [2.0, 4.0, 6.0, 8.0])
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|