# 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_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()