diff --git a/test/imported/test_indexing.py b/test/imported/test_indexing.py index abf804d359..d7d543987e 100644 --- a/test/imported/test_indexing.py +++ b/test/imported/test_indexing.py @@ -1,9 +1,11 @@ # test cases are modified from pytorch test_indexing.py https://github.com/pytorch/pytorch/blob/597d3fb86a2f3b8d6d8ee067e769624dcca31cdb/test/test_indexing.py -import math, unittest, random +import math, unittest, random, copy +# import warnings import numpy as np -from tinygrad.tensor import Tensor, dtypes +from tinygrad import Tensor, dtypes +# from tinygrad import TinyJit from tinygrad.lazy import LazyBuffer from tinygrad.shape.shapetracker import ShapeTracker from tinygrad.shape.view import View @@ -27,7 +29,42 @@ def set_(reference: Tensor, shape, strides, offset): assert strided.lazydata in reference.lazydata.base.views, "base.views should contain strided.lazydata" return strided -def make_tensor(shape, dtype:dtypes, noncontiguous): +# TODO tries to mimic .detach().copy() or just .copy() behavior +# for torch.copy() the resulting tensor.data_ptr() is different than the original +# some random medium article says .clone() is a shallow copy, so maybe this is correct? +# https://discuss.pytorch.org/t/clone-and-detach-in-v0-4-0/16861/4 +def clone(original:Tensor): + ret = Tensor(copy.copy(original.lazydata), device=original.device, dtype=original.dtype, requires_grad=original.requires_grad) + # TODO nah this isn't right either. maybe compare their loadops??? idk + assert ret.lazydata is not original.lazydata, "clone should create a lazydata copy" + return ret + +# TODO torch.data_ptr +def data_ptr(tensor:Tensor): + ... + +# TODO torch.argsort() +def argsort(tensor:Tensor) -> Tensor: + ... + +# TODO torch.all() +# not sure if this is right +def all_(tensor:Tensor) -> Tensor: + return tensor != 0 + +def diagonal(tensor:Tensor) -> Tensor: + assert all(sh == sh[0] for sh in tensor.shape), "matrix should be square" + assert tensor.ndim == 2, 'only support 2 ndim tensors' + return (Tensor.eye(tensor.shape[0]) * tensor).sum(0) + +# TODO torch.copy_() +def copy_(src:Tensor, other:Tensor) -> Tensor: + ... + +def unravel_index(tensor, shape): + ... + +def make_tensor(shape, dtype:dtypes, noncontiguous) -> Tensor: r"""Creates a tensor with the given :attr:`shape`, :attr:`device`, and :attr:`dtype`, and filled with values uniformly drawn from ``[low, high)``. @@ -51,9 +88,9 @@ def make_tensor(shape, dtype:dtypes, noncontiguous): +---------------------------+------------+----------+ """ contiguous = not noncontiguous # lol - if dtype is dtypes.bool: return Tensor.randint(shape=shape, low=0, high=2, dtype=dtypes.bool, contiguous=contiguous) - elif dtype.is_unsigned(): return Tensor.randint(shape=shape, low=0, high=10, dtype=dtype, contiguous=contiguous) - elif dtype.is_int(): return Tensor.randint(shape=shape, low=-9, high=10, dtype=dtype, contiguous=contiguous) # signed int + if dtype is dtypes.bool: return Tensor.randint(shape=shape, low=0, high=2, contiguous=contiguous).cast(dtypes.bool) + elif dtype.is_unsigned(): return Tensor.randint(shape=shape, low=0, high=10, contiguous=contiguous).cast(dtype) + elif dtype.is_int(): return Tensor.randint(shape=shape, low=-9, high=10, contiguous=contiguous).cast(dtype) # signed int elif dtype.is_float(): return Tensor.rand(shape=shape, low=-9, high=9, dtype=dtype, contiguous=contiguous) else: raise NotImplementedError(f"{dtype} not implemented") @@ -156,7 +193,7 @@ class TestIndexing(unittest.TestCase): # pick a random valid indexer type def ri(indices): choice = random.randint(0, 2) - if choice == 0: return Tensor(indices) + if choice == 0: return Tensor(indices, dtype=dtypes.int32) if choice == 1: return list(indices) return tuple(indices) @@ -364,7 +401,7 @@ class TestIndexing(unittest.TestCase): ''' strided[ri([0]), ri([1])] = -1 numpy_testing_assert_equal_helper(strided[ri([0]), ri([1])], - np.array([-1])) + Tensor([-1])) ''' reference = Tensor.arange(0., 24).reshape(3, 8) @@ -374,9 +411,9 @@ class TestIndexing(unittest.TestCase): np.array([11, 17])) # TODO setitem ''' - strided[ri([0, 1]), ri([1, 0])] = np.array([-1, 2]) + strided[ri([0, 1]), ri([1, 0])] = Tensor([-1, 2]) numpy_testing_assert_equal_helper(strided[ri([0, 1]), ri([1, 0])], - np.array([-1, 2])) + Tensor([-1, 2])) ''' reference = Tensor.arange(0., 24).realize().reshape(3, 8) @@ -390,9 +427,9 @@ class TestIndexing(unittest.TestCase): np.array([[10, 11], [17, 18]])) # TODO setitem ''' - strided[rows, columns] = np.array([[4, 6], [2, 3]]) + strided[rows, columns] = Tensor([[4, 6], [2, 3]]) numpy_testing_assert_equal_helper(strided[rows, columns], - np.array([[4, 6], [2, 3]])) + Tensor([[4, 6], [2, 3]])) ''' # Tests using less than the number of dims, and ellipsis @@ -461,15 +498,15 @@ class TestIndexing(unittest.TestCase): numpy_testing_assert_equal_helper(pyt, numt) ''' def assert_set_eq(tensor: Tensor, indexer, val): - pyt = tensor.detach() - numt = tensor.detach() + pyt = clone(tensor) + numt = clone(tensor) pyt[indexer] = val numt = set_numpy(numt, indexer, val) numpy_testing_assert_equal_helper(pyt, numt) # NOTE: torch initiates the gradients using g0cpu (rand as gradients) def assert_backward_eq(tensor: Tensor, indexer): - cpu = tensor.float().detach() + cpu = clone(tensor.float()) cpu.requires_grad = True outcpu = cpu[indexer].sum() outcpu.backward() @@ -667,15 +704,18 @@ class TestIndexing(unittest.TestCase): ''' assert_backward_eq(reference, indexer) - # def test_set_item_to_scalar_tensor(self): - # m = random.randint(1, 10) - # n = random.randint(1, 10) - # z = torch.randn([m, n]) - # a = 1.0 - # w = np.array(a, requires_grad=True) - # z[:, 0] = w - # z.sum().backward() - # numpy_testing_assert_equal_helper(w.grad, m * a) + # TODO setitem + ''' + def test_set_item_to_scalar_tensor(self): + m = random.randint(1, 10) + n = random.randint(1, 10) + z = Tensor.randn([m, n]) + a = 1.0 + w = Tensor(a, requires_grad=True) + z[:, 0] = w + z.sum().backward() + numpy_testing_assert_equal_helper(w.grad, m * a) + ''' def test_single_int(self): v = Tensor.randn(5, 7, 3) @@ -701,575 +741,688 @@ class TestIndexing(unittest.TestCase): numpy_testing_assert_equal_helper(v[::11], [0]) numpy_testing_assert_equal_helper(v[1:6:2], [1, 3, 5]) - # def test_step_assignment(self): - # v = torch.zeros(4, 4) - # v[0, 1::2] = np.array([3., 4.]) - # numpy_testing_assert_equal_helper(v[0].tolist(), [0, 3, 0, 4]) - # numpy_testing_assert_equal_helper(v[1:].sum(), 0) - - # def test_bool_indices(self): - # v = Tensor.randn(5, 7, 3) - # boolIndices = np.array([True, False, True, True, False], dtype=bool) - # numpy_testing_assert_equal_helper(v[boolIndices].shape, (3, 7, 3)) - # numpy_testing_assert_equal_helper(v[boolIndices], Tensor.stack([v[0], v[2], v[3]])) - - # v = np.array([True, False, True], dtype=torch.bool) - # boolIndices = np.array([True, False, False], dtype=torch.bool) - # uint8Indices = np.array([1, 0, 0], dtype=torch.uint8) - # with warnings.catch_warnings(record=True) as w: - # numpy_testing_assert_equal_helper(v[boolIndices].shape, v[uint8Indices].shape) - # numpy_testing_assert_equal_helper(v[boolIndices], v[uint8Indices]) - # numpy_testing_assert_equal_helper(v[boolIndices], tensor([True], dtype=torch.bool)) - # numpy_testing_assert_equal_helper(len(w), 2) - - # def test_bool_indices_accumulate(self): - # mask = torch.zeros(size=(10, ), dtype=torch.bool) - # y = torch.ones(size=(10, 10)) - # y.index_put_((mask, ), y[mask], accumulate=True) - # numpy_testing_assert_equal_helper(y, torch.ones(size=(10, 10))) - - # def test_multiple_bool_indices(self): - # v = torch.randn(5, 7, 3) - # # note: these broadcast together and are transposed to the first dim - # mask1 = np.array([1, 0, 1, 1, 0], dtype=torch.bool) - # mask2 = np.array([1, 1, 1], dtype=torch.bool) - # numpy_testing_assert_equal_helper(v[mask1, :, mask2].shape, (3, 7)) - - # def test_byte_mask(self): - # v = torch.randn(5, 7, 3) - # mask = torch.ByteTensor([1, 0, 1, 1, 0]).to(device) - # with warnings.catch_warnings(record=True) as w: - # numpy_testing_assert_equal_helper(v[mask].shape, (3, 7, 3)) - # numpy_testing_assert_equal_helper(v[mask], torch.stack([v[0], v[2], v[3]])) - # numpy_testing_assert_equal_helper(len(w), 2) - - # v = np.array([1.]) - # numpy_testing_assert_equal_helper(v[v == 0], np.array([])) - - # def test_byte_mask_accumulate(self): - # mask = torch.zeros(size=(10, ), dtype=torch.uint8) - # y = torch.ones(size=(10, 10)) - # with warnings.catch_warnings(record=True) as w: - # warnings.simplefilter("always") - # y.index_put_((mask, ), y[mask], accumulate=True) - # numpy_testing_assert_equal_helper(y, torch.ones(size=(10, 10))) - # numpy_testing_assert_equal_helper(len(w), 2) - - # def test_index_put_accumulate_large_tensor(self): - # # This test is for tensors with number of elements >= INT_MAX (2^31 - 1). - # N = (1 << 31) + 5 - # dt = torch.int8 - # a = torch.ones(N, dtype=dt) - # indices = np.array([-2, 0, -2, -1, 0, -1, 1], dtype=torch.long) - # values = np.array([6, 5, 6, 6, 5, 7, 11], dtype=dt) - - # a.index_put_((indices, ), values, accumulate=True) - - # numpy_testing_assert_equal_helper(a[0], 11) - # numpy_testing_assert_equal_helper(a[1], 12) - # numpy_testing_assert_equal_helper(a[2], 1) - # numpy_testing_assert_equal_helper(a[-3], 1) - # numpy_testing_assert_equal_helper(a[-2], 13) - # numpy_testing_assert_equal_helper(a[-1], 14) - - # a = torch.ones((2, N), dtype=dt) - # indices0 = np.array([0, -1, 0, 1], dtype=torch.long) - # indices1 = np.array([-2, -1, 0, 1], dtype=torch.long) - # values = np.array([12, 13, 10, 11], dtype=dt) - - # a.index_put_((indices0, indices1), values, accumulate=True) - - # numpy_testing_assert_equal_helper(a[0, 0], 11) - # numpy_testing_assert_equal_helper(a[0, 1], 1) - # numpy_testing_assert_equal_helper(a[1, 0], 1) - # numpy_testing_assert_equal_helper(a[1, 1], 12) - # numpy_testing_assert_equal_helper(a[:, 2], torch.ones(2, dtype=torch.int8)) - # numpy_testing_assert_equal_helper(a[:, -3], torch.ones(2, dtype=torch.int8)) - # numpy_testing_assert_equal_helper(a[0, -2], 13) - # numpy_testing_assert_equal_helper(a[1, -2], 1) - # numpy_testing_assert_equal_helper(a[-1, -1], 14) - # numpy_testing_assert_equal_helper(a[0, -1], 1) - - # def test_index_put_accumulate_expanded_values(self): - # # checks the issue with cuda: https://github.com/pytorch/pytorch/issues/39227 - # # and verifies consistency with CPU result - # t = torch.zeros((5, 2)) - # t_dev = t.to(device) - # indices = [ - # np.array([0, 1, 2, 3]), - # np.array([1, ]), - # ] - # indices_dev = [i.to(device) for i in indices] - # values0d = np.array(1.0) - # values1d = np.array([1.0, ]) - - # out_cuda = t_dev.index_put_(indices_dev, values0d.to(device), accumulate=True) - # out_cpu = t.index_put_(indices, values0d, accumulate=True) - # numpy_testing_assert_equal_helper(out_cuda.cpu(), out_cpu) - - # out_cuda = t_dev.index_put_(indices_dev, values1d.to(device), accumulate=True) - # out_cpu = t.index_put_(indices, values1d, accumulate=True) - # numpy_testing_assert_equal_helper(out_cuda.cpu(), out_cpu) - - # t = torch.zeros(4, 3, 2) - # t_dev = t.to(device) - - # indices = [ - # np.array([0, ]), - # torch.arange(3)[:, None], - # torch.arange(2)[None, :], - # ] - # indices_dev = [i.to(device) for i in indices] - # values1d = np.array([-1.0, -2.0]) - # values2d = np.array([[-1.0, -2.0], ]) - - # out_cuda = t_dev.index_put_(indices_dev, values1d.to(device), accumulate=True) - # out_cpu = t.index_put_(indices, values1d, accumulate=True) - # numpy_testing_assert_equal_helper(out_cuda.cpu(), out_cpu) - - # out_cuda = t_dev.index_put_(indices_dev, values2d.to(device), accumulate=True) - # out_cpu = t.index_put_(indices, values2d, accumulate=True) - # numpy_testing_assert_equal_helper(out_cuda.cpu(), out_cpu) - - # def test_index_put_accumulate_non_contiguous(self): - # t = torch.zeros((5, 2, 2)) - # t_dev = t.to(device) - # t1 = t_dev[:, 0, :] - # t2 = t[:, 0, :] - # self.assertTrue(not t1.is_contiguous()) - # self.assertTrue(not t2.is_contiguous()) - - # indices = [np.array([0, 1]), ] - # indices_dev = [i.to(device) for i in indices] - # value = torch.randn(2, 2) - # out_cuda = t1.index_put_(indices_dev, value.to(device), accumulate=True) - # out_cpu = t2.index_put_(indices, value, accumulate=True) - # self.assertTrue(not t1.is_contiguous()) - # self.assertTrue(not t2.is_contiguous()) - - # numpy_testing_assert_equal_helper(out_cuda.cpu(), out_cpu) - - # def test_index_put_accumulate_with_optional_tensors(self): - # # TODO: replace with a better solution. - # # Currently, here using torchscript to put None into indices. - # # on C++ it gives indices as a list of 2 optional tensors: first is null and - # # the second is a valid tensor. - # @torch.jit.script - # def func(x, i, v): - # idx = [None, i] - # x.index_put_(idx, v, accumulate=True) - # return x - - # n = 4 - # t = torch.arange(n * 2, dtype=torch.float32).reshape(n, 2) - # t_dev = t.to(device) - # indices = np.array([1, 0]) - # indices_dev = indices.to(device) - # value0d = np.array(10.0) - # value1d = np.array([1.0, 2.0]) - - # out_cuda = func(t_dev, indices_dev, value0d.cuda()) - # out_cpu = func(t, indices, value0d) - # numpy_testing_assert_equal_helper(out_cuda.cpu(), out_cpu) - - # out_cuda = func(t_dev, indices_dev, value1d.cuda()) - # out_cpu = func(t, indices, value1d) - # numpy_testing_assert_equal_helper(out_cuda.cpu(), out_cpu) - - # def test_index_put_accumulate_duplicate_indices(self): - # for i in range(1, 512): - # # generate indices by random walk, this will create indices with - # # lots of duplicates interleaved with each other - # delta = torch.empty(i, dtype=torch.double).uniform_(-1, 1) - # indices = delta.cumsum(0).long() - - # input = torch.randn(indices.abs().max() + 1) - # values = torch.randn(indices.size(0)) - # output = input.index_put((indices,), values, accumulate=True) - - # input_list = input.tolist() - # indices_list = indices.tolist() - # values_list = values.tolist() - # for i, v in zip(indices_list, values_list): - # input_list[i] += v - - # numpy_testing_assert_equal_helper(output, input_list) - - # def test_index_ind_dtype(self): - # x = torch.randn(4, 4) - # ind_long = torch.randint(4, (4,), dtype=torch.long) - # ind_int = ind_long.int() - # src = torch.randn(4) - # ref = x[ind_long, ind_long] - # res = x[ind_int, ind_int] - # numpy_testing_assert_equal_helper(ref, res) - # ref = x[ind_long, :] - # res = x[ind_int, :] - # numpy_testing_assert_equal_helper(ref, res) - # ref = x[:, ind_long] - # res = x[:, ind_int] - # numpy_testing_assert_equal_helper(ref, res) - # # no repeating indices for index_put - # ind_long = torch.arange(4, dtype=torch.long) - # ind_int = ind_long.int() - # for accum in (True, False): - # inp_ref = x.clone() - # inp_res = x.clone() - # torch.index_put_(inp_ref, (ind_long, ind_long), src, accum) - # torch.index_put_(inp_res, (ind_int, ind_int), src, accum) - # numpy_testing_assert_equal_helper(inp_ref, inp_res) - - # def test_index_put_accumulate_empty(self): - # # Regression test for https://github.com/pytorch/pytorch/issues/94667 - # input = torch.rand([], dtype=torch.float32) - # with self.assertRaises(RuntimeError): - # input.index_put([], np.array([1.0]), True) - - # def test_multiple_byte_mask(self): - # v = torch.randn(5, 7, 3) - # # note: these broadcast together and are transposed to the first dim - # mask1 = torch.ByteTensor([1, 0, 1, 1, 0]).to(device) - # mask2 = torch.ByteTensor([1, 1, 1]).to(device) - # with warnings.catch_warnings(record=True) as w: - # warnings.simplefilter("always") - # numpy_testing_assert_equal_helper(v[mask1, :, mask2].shape, (3, 7)) - # numpy_testing_assert_equal_helper(len(w), 2) - - # def test_byte_mask2d(self): - # v = torch.randn(5, 7, 3) - # c = torch.randn(5, 7) - # num_ones = (c > 0).sum() - # r = v[c > 0] - # numpy_testing_assert_equal_helper(r.shape, (num_ones, 3)) - - # def test_jit_indexing(self): - # def fn1(x): - # x[x < 50] = 1.0 - # return x - - # def fn2(x): - # x[0:50] = 1.0 - # return x - - # scripted_fn1 = torch.jit.script(fn1) - # scripted_fn2 = torch.jit.script(fn2) - # data = torch.arange(100, dtype=torch.float) - # out = scripted_fn1(data.detach().clone()) - # ref = np.array(np.concatenate((np.ones(50), np.arange(50, 100))), dtype=torch.float) - # numpy_testing_assert_equal_helper(out, ref) - # out = scripted_fn2(data.detach().clone()) - # numpy_testing_assert_equal_helper(out, ref) - - # def test_int_indices(self): - # v = torch.randn(5, 7, 3) - # numpy_testing_assert_equal_helper(v[[0, 4, 2]].shape, (3, 7, 3)) - # numpy_testing_assert_equal_helper(v[:, [0, 4, 2]].shape, (5, 3, 3)) - # numpy_testing_assert_equal_helper(v[:, [[0, 1], [4, 3]]].shape, (5, 2, 2, 3)) - - # def test_index_put_src_datatype(self, dtype): - # src = torch.ones(3, 2, 4, dtype=dtype) - # vals = torch.ones(3, 2, 4, dtype=dtype) - # indices = (np.array([0, 2, 1]),) - # res = src.index_put_(indices, vals, accumulate=True) - # numpy_testing_assert_equal_helper(res.shape, src.shape) - - # def test_index_src_datatype(self, dtype): - # src = torch.ones(3, 2, 4, dtype=dtype) - # # test index - # res = src[[0, 2, 1], :, :] - # numpy_testing_assert_equal_helper(res.shape, src.shape) - # # test index_put, no accum - # src[[0, 2, 1], :, :] = res - # numpy_testing_assert_equal_helper(res.shape, src.shape) - - # def test_int_indices2d(self): - # # From the NumPy indexing example - # x = torch.arange(0, 12).view(4, 3) - # rows = np.array([[0, 0], [3, 3]]) - # columns = np.array([[0, 2], [0, 2]]) - # numpy_testing_assert_equal_helper(x[rows, columns].tolist(), [[0, 2], [9, 11]]) - - # def test_int_indices_broadcast(self): - # # From the NumPy indexing example - # x = torch.arange(0, 12).view(4, 3) - # rows = np.array([0, 3]) - # columns = np.array([0, 2]) - # result = x[rows[:, None], columns] - # numpy_testing_assert_equal_helper(result.tolist(), [[0, 2], [9, 11]]) - - # def test_empty_index(self): - # x = torch.arange(0, 12).view(4, 3) - # idx = np.array([], dtype=torch.long) - # numpy_testing_assert_equal_helper(x[idx].numel(), 0) - - # # empty assignment should have no effect but not throw an exception - # y = x.clone() - # y[idx] = -1 - # numpy_testing_assert_equal_helper(x, y) - - # mask = torch.zeros(4, 3).bool() - # y[mask] = -1 - # numpy_testing_assert_equal_helper(x, y) - - # def test_empty_ndim_index(self): - # x = torch.randn(5) - # numpy_testing_assert_equal_helper(torch.empty(0, 2), x[torch.empty(0, 2, dtype=torch.int64)]) - - # x = torch.randn(2, 3, 4, 5) - # numpy_testing_assert_equal_helper(torch.empty(2, 0, 6, 4, 5), - # x[:, torch.empty(0, 6, dtype=torch.int64)]) - - # x = torch.empty(10, 0) - # numpy_testing_assert_equal_helper(x[[1, 2]].shape, (2, 0)) - # numpy_testing_assert_equal_helper(x[[], []].shape, (0,)) - # with self.assertRaisesRegex(IndexError, 'for dimension with size 0'): - # x[:, [0, 1]] - - # def test_empty_ndim_index_bool(self): - # x = torch.randn(5) - # self.assertRaises(IndexError, lambda: x[torch.empty(0, 2, dtype=torch.uint8)]) - - # def test_empty_slice(self): - # x = torch.randn(2, 3, 4, 5) - # y = x[:, :, :, 1] - # z = y[:, 1:1, :] - # numpy_testing_assert_equal_helper((2, 0, 4), z.shape) - # # this isn't technically necessary, but matches NumPy stride calculations. - # numpy_testing_assert_equal_helper((60, 20, 5), z.stride()) - # self.assertTrue(z.is_contiguous()) - - # def test_index_getitem_copy_bools_slices(self): - # true = np.array(1, dtype=torch.uint8) - # false = np.array(0, dtype=torch.uint8) - - # tensors = [torch.randn(2, 3), np.array(3.)] - - # for a in tensors: - # self.assertNotEqual(a.data_ptr(), a[True].data_ptr()) - # numpy_testing_assert_equal_helper(torch.empty(0, *a.shape), a[False]) - # self.assertNotEqual(a.data_ptr(), a[true].data_ptr()) - # numpy_testing_assert_equal_helper(torch.empty(0, *a.shape), a[false]) - # numpy_testing_assert_equal_helper(a.data_ptr(), a[None].data_ptr()) - # numpy_testing_assert_equal_helper(a.data_ptr(), a[...].data_ptr()) - - # def test_index_setitem_bools_slices(self): - # true = np.array(1, dtype=torch.uint8) - # false = np.array(0, dtype=torch.uint8) - - # tensors = [torch.randn(2, 3), np.array(3)] - - # for a in tensors: - # # prefix with a 1,1, to ensure we are compatible with numpy which cuts off prefix 1s - # # (some of these ops already prefix a 1 to the size) - # neg_ones = torch.ones_like(a) * -1 - # neg_ones_expanded = neg_ones.unsqueeze(0).unsqueeze(0) - # a[True] = neg_ones_expanded - # numpy_testing_assert_equal_helper(a, neg_ones) - # a[False] = 5 - # numpy_testing_assert_equal_helper(a, neg_ones) - # a[true] = neg_ones_expanded * 2 - # numpy_testing_assert_equal_helper(a, neg_ones * 2) - # a[false] = 5 - # numpy_testing_assert_equal_helper(a, neg_ones * 2) - # a[None] = neg_ones_expanded * 3 - # numpy_testing_assert_equal_helper(a, neg_ones * 3) - # a[...] = neg_ones_expanded * 4 - # numpy_testing_assert_equal_helper(a, neg_ones * 4) - # if a.dim() == 0: - # with self.assertRaises(IndexError): - # a[:] = neg_ones_expanded * 5 - - # def test_index_scalar_with_bool_mask(self): - # a = np.array(1) - # uintMask = np.array(True, dtype=torch.uint8) - # boolMask = np.array(True, dtype=torch.bool) - # numpy_testing_assert_equal_helper(a[uintMask], a[boolMask]) - # numpy_testing_assert_equal_helper(a[uintMask].dtype, a[boolMask].dtype) - - # a = np.array(True, dtype=torch.bool) - # numpy_testing_assert_equal_helper(a[uintMask], a[boolMask]) - # numpy_testing_assert_equal_helper(a[uintMask].dtype, a[boolMask].dtype) - - # def test_setitem_expansion_error(self): - # true = np.array(True) - # a = torch.randn(2, 3) - # # check prefix with non-1s doesn't work - # a_expanded = a.expand(torch.Size([5, 1]) + a.size()) - # # NumPy: ValueError - # with self.assertRaises(RuntimeError): - # a[True] = a_expanded - # with self.assertRaises(RuntimeError): - # a[true] = a_expanded - - # def test_getitem_scalars(self): - # zero = np.array(0, dtype=torch.int64) - # one = np.array(1, dtype=torch.int64) - - # # non-scalar indexed with scalars - # a = torch.randn(2, 3) - # numpy_testing_assert_equal_helper(a[0], a[zero]) - # numpy_testing_assert_equal_helper(a[0][1], a[zero][one]) - # numpy_testing_assert_equal_helper(a[0, 1], a[zero, one]) - # numpy_testing_assert_equal_helper(a[0, one], a[zero, 1]) - - # # indexing by a scalar should slice (not copy) - # numpy_testing_assert_equal_helper(a[0, 1].data_ptr(), a[zero, one].data_ptr()) - # numpy_testing_assert_equal_helper(a[1].data_ptr(), a[one.int()].data_ptr()) - # numpy_testing_assert_equal_helper(a[1].data_ptr(), a[one.short()].data_ptr()) - - # # scalar indexed with scalar - # r = torch.randn(()) - # with self.assertRaises(IndexError): - # r[:] - # with self.assertRaises(IndexError): - # r[zero] - # numpy_testing_assert_equal_helper(r, r[...]) - - # def test_setitem_scalars(self): - # zero = np.array(0, dtype=torch.int64) - - # # non-scalar indexed with scalars - # a = torch.randn(2, 3) - # a_set_with_number = a.clone() - # a_set_with_scalar = a.clone() - # b = torch.randn(3) - - # a_set_with_number[0] = b - # a_set_with_scalar[zero] = b - # numpy_testing_assert_equal_helper(a_set_with_number, a_set_with_scalar) - # a[1, zero] = 7.7 - # numpy_testing_assert_equal_helper(7.7, a[1, 0]) - - # # scalar indexed with scalars - # r = torch.randn(()) - # with self.assertRaises(IndexError): - # r[:] = 8.8 - # with self.assertRaises(IndexError): - # r[zero] = 8.8 - # r[...] = 9.9 - # numpy_testing_assert_equal_helper(9.9, r) - - # def test_basic_advanced_combined(self): - # # From the NumPy indexing example - # x = torch.arange(0, 12).view(4, 3) - # numpy_testing_assert_equal_helper(x[1:2, 1:3], x[1:2, [1, 2]]) - # numpy_testing_assert_equal_helper(x[1:2, 1:3].tolist(), [[4, 5]]) - - # # Check that it is a copy - # unmodified = x.clone() - # x[1:2, [1, 2]].zero_() - # numpy_testing_assert_equal_helper(x, unmodified) - - # # But assignment should modify the original - # unmodified = x.clone() - # x[1:2, [1, 2]] = 0 - # self.assertNotEqual(x, unmodified) - - # def test_int_assignment(self): - # x = torch.arange(0, 4).view(2, 2) - # x[1] = 5 - # numpy_testing_assert_equal_helper(x.tolist(), [[0, 1], [5, 5]]) - - # x = torch.arange(0, 4).view(2, 2) - # x[1] = torch.arange(5, 7) - # numpy_testing_assert_equal_helper(x.tolist(), [[0, 1], [5, 6]]) - - # def test_byte_tensor_assignment(self): - # x = torch.arange(0., 16).view(4, 4) - # b = torch.ByteTensor([True, False, True, False]).to(device) - # value = np.array([3., 4., 5., 6.]) - - # with warnings.catch_warnings(record=True) as w: - # x[b] = value - # numpy_testing_assert_equal_helper(len(w), 1) - - # numpy_testing_assert_equal_helper(x[0], value) - # numpy_testing_assert_equal_helper(x[1], torch.arange(4., 8)) - # numpy_testing_assert_equal_helper(x[2], value) - # numpy_testing_assert_equal_helper(x[3], torch.arange(12., 16)) - - # def test_variable_slicing(self): - # x = torch.arange(0, 16).view(4, 4) - # indices = torch.IntTensor([0, 1]).to(device) - # i, j = indices - # numpy_testing_assert_equal_helper(x[i:j], x[0:1]) - - # def test_ellipsis_tensor(self): - # x = torch.arange(0, 9).view(3, 3) - # idx = np.array([0, 2]) - # numpy_testing_assert_equal_helper(x[..., idx].tolist(), [[0, 2], - # [3, 5], - # [6, 8]]) - # numpy_testing_assert_equal_helper(x[idx, ...].tolist(), [[0, 1, 2], - # [6, 7, 8]]) - + # TODO setitem + ''' + def test_step_assignment(self): + v = Tensor.zeros(4, 4) + v[0, 1::2] = Tensor([3., 4.]) + numpy_testing_assert_equal_helper(v[0].numpy().tolist(), [0, 3, 0, 4]) + numpy_testing_assert_equal_helper(v[1:].sum(), 0) + ''' + + # TODO bool indexing + ''' + def test_bool_indices(self): + v = Tensor.randn(5, 7, 3) + boolIndices = Tensor([True, False, True, True, False], dtype=dtypes.bool) + numpy_testing_assert_equal_helper(v[boolIndices].shape, (3, 7, 3)) + numpy_testing_assert_equal_helper(v[boolIndices], Tensor.stack([v[0], v[2], v[3]])) + + v = Tensor([True, False, True], dtype=dtypes.bool) + boolIndices = Tensor([True, False, False], dtype=dtypes.bool) + uint8Indices = Tensor([1, 0, 0], dtype=dtypes.uint8) + with warnings.catch_warnings(record=True) as w: + numpy_testing_assert_equal_helper(v[boolIndices].shape, v[uint8Indices].shape) + numpy_testing_assert_equal_helper(v[boolIndices], v[uint8Indices]) + numpy_testing_assert_equal_helper(v[boolIndices], tensor([True], dtype=torch.bool)) + numpy_testing_assert_equal_helper(len(w), 2) + ''' + + # TODO setindex + # TODO bool indexing + ''' + def test_bool_indices_accumulate(self): + mask = Tensor.zeros(size=(10, ), dtype=dtypes.bool) + y = Tensor.ones(size=(10, 10)) + y.index_put_((mask, ), y[mask], accumulate=True) + numpy_testing_assert_equal_helper(y, Tensor.ones(size=(10, 10))) + ''' + + # TODO bool indexing + ''' + def test_multiple_bool_indices(self): + v = Tensor.randn(5, 7, 3) + # note: these broadcast together and are transposed to the first dim + mask1 = Tensor([1, 0, 1, 1, 0], dtype=dtypes.bool) + mask2 = Tensor([1, 1, 1], dtype=dtypes.bool) + numpy_testing_assert_equal_helper(v[mask1, :, mask2].shape, (3, 7)) + ''' + + # TODO bool indexing + ''' + def test_byte_mask(self): + v = Tensor.randn(5, 7, 3) + mask = Tensor([1, 0, 1, 1, 0], dtype=dtypes.uint8) + with warnings.catch_warnings(record=True) as w: + numpy_testing_assert_equal_helper(v[mask].shape, (3, 7, 3)) + numpy_testing_assert_equal_helper(v[mask], Tensor.stack([v[0], v[2], v[3]])) + numpy_testing_assert_equal_helper(len(w), 2) + + v = Tensor([1.]) + numpy_testing_assert_equal_helper(v[v == 0], Tensor([])) + ''' + + # TODO setitem + # TODO bool indexing + ''' + def test_byte_mask_accumulate(self): + mask = Tensor.zeros(size=(10, ), dtype=dtypes.uint8) + y = Tensor.ones(size=(10, 10)) + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + y.index_put_((mask, ), y[mask], accumulate=True) + numpy_testing_assert_equal_helper(y, Tensor.ones(size=(10, 10))) + numpy_testing_assert_equal_helper(len(w), 2) + ''' + + # TODO setitem + ''' + def test_index_put_accumulate_large_tensor(self): + # This test is for tensors with number of elements >= INT_MAX (2^31 - 1). + N = (1 << 31) + 5 + dt = dtypes.int8 + a = Tensor.ones(N, dtype=dt) + indices = Tensor([-2, 0, -2, -1, 0, -1, 1], dtype=dtypes.int64) + values = Tensor([6, 5, 6, 6, 5, 7, 11], dtype=dt) + + a.index_put_((indices, ), values, accumulate=True) + + numpy_testing_assert_equal_helper(a[0], 11) + numpy_testing_assert_equal_helper(a[1], 12) + numpy_testing_assert_equal_helper(a[2], 1) + numpy_testing_assert_equal_helper(a[-3], 1) + numpy_testing_assert_equal_helper(a[-2], 13) + numpy_testing_assert_equal_helper(a[-1], 14) + + a = Tensor.ones((2, N), dtype=dt) + indices0 = np.array([0, -1, 0, 1], dtype=dtypes.int64) + indices1 = np.array([-2, -1, 0, 1], dtype=dtypes.int64) + values = np.array([12, 13, 10, 11], dtype=dt) + + a.index_put_((indices0, indices1), values, accumulate=True) + + numpy_testing_assert_equal_helper(a[0, 0], 11) + numpy_testing_assert_equal_helper(a[0, 1], 1) + numpy_testing_assert_equal_helper(a[1, 0], 1) + numpy_testing_assert_equal_helper(a[1, 1], 12) + numpy_testing_assert_equal_helper(a[:, 2], Tensor.ones(2, dtype=dtypes.int8)) + numpy_testing_assert_equal_helper(a[:, -3], Tensor.ones(2, dtype=dtypes.int8)) + numpy_testing_assert_equal_helper(a[0, -2], 13) + numpy_testing_assert_equal_helper(a[1, -2], 1) + numpy_testing_assert_equal_helper(a[-1, -1], 14) + numpy_testing_assert_equal_helper(a[0, -1], 1) + ''' + + @unittest.skip("pytorch device specific test") + def test_index_put_accumulate_expanded_values(self, device): + # checks the issue with cuda: https://github.com/pytorch/pytorch/issues/39227 + # and verifies consistency with CPU result + t = Tensor.zeros((5, 2)) + t_dev = t.to(device) + indices = [ + Tensor([0, 1, 2, 3]), + Tensor([1, ]), + ] + indices_dev = [i.to(device) for i in indices] + values0d = Tensor(1.0) + values1d = Tensor([1.0, ]) + + out_cuda = t_dev.index_put_(indices_dev, values0d.to(device), accumulate=True) + out_cpu = t.index_put_(indices, values0d, accumulate=True) + numpy_testing_assert_equal_helper(out_cuda.cpu(), out_cpu) + + out_cuda = t_dev.index_put_(indices_dev, values1d.to(device), accumulate=True) + out_cpu = t.index_put_(indices, values1d, accumulate=True) + numpy_testing_assert_equal_helper(out_cuda.cpu(), out_cpu) + + t = Tensor.zeros(4, 3, 2) + t_dev = t.to(device) + + indices = [ + Tensor([0, ]), + Tensor.arange(3)[:, None], + Tensor.arange(2)[None, :], + ] + indices_dev = [i.to(device) for i in indices] + values1d = np.array([-1.0, -2.0]) + values2d = np.array([[-1.0, -2.0], ]) + + out_cuda = t_dev.index_put_(indices_dev, values1d.to(device), accumulate=True) + out_cpu = t.index_put_(indices, values1d, accumulate=True) + numpy_testing_assert_equal_helper(out_cuda.cpu(), out_cpu) + + out_cuda = t_dev.index_put_(indices_dev, values2d.to(device), accumulate=True) + out_cpu = t.index_put_(indices, values2d, accumulate=True) + numpy_testing_assert_equal_helper(out_cuda.cpu(), out_cpu) + + # TODO setitem + @unittest.skip("pytorch device specific test") + def test_index_put_accumulate_non_contiguous(self, device): + t = Tensor.zeros((5, 2, 2)) + t_dev = t.to(device) + t1 = t_dev[:, 0, :] + t2 = t[:, 0, :] + self.assertTrue(not t1.lazydata.st.contiguous) + self.assertTrue(not t2.lazydata.st.contiguous) + + indices = [Tensor([0, 1]), ] + indices_dev = [i.to(device) for i in indices] + value = Tensor.randn(2, 2) + out_cuda = t1.index_put_(indices_dev, value.to(device), accumulate=True) + out_cpu = t2.index_put_(indices, value, accumulate=True) + self.assertTrue(not t1.lazydata.st.contiguous) + self.assertTrue(not t2.lazydata.st.contiguous) + + numpy_testing_assert_equal_helper(out_cuda.cpu(), out_cpu) + + # TODO setitem + ''' + def test_index_put_accumulate_with_optional_tensors(self): + # TODO: replace with a better solution. + # Currently, here using torchscript to put None into indices. + # on C++ it gives indices as a list of 2 optional tensors: first is null and + # the second is a valid tensor. + @TinyJit + def func(x, i, v): + idx = [None, i] + x.index_put_(idx, v, accumulate=True) + return x + + n = 4 + t = Tensor.arange(n * 2, dtype=dtypes.float32).reshape(n, 2) + t_dev = t.to(device) + indices = Tensor([1, 0]) + indices_dev = indices.to(device) + value0d = Tensor(10.0) + value1d = Tensor([1.0, 2.0]) + + out_cuda = func(t_dev, indices_dev, value0d.cuda()) + out_cpu = func(t, indices, value0d) + numpy_testing_assert_equal_helper(out_cuda.cpu(), out_cpu) + + out_cuda = func(t_dev, indices_dev, value1d.cuda()) + out_cpu = func(t, indices, value1d) + numpy_testing_assert_equal_helper(out_cuda.cpu(), out_cpu) + ''' + + # TODO setindex + ''' + def test_index_put_accumulate_duplicate_indices(self): + for i in range(1, 512): + # generate indices by random walk, this will create indices with + # lots of duplicates interleaved with each other + delta = Tensor.uniform(low=-1, high=1, dtype=dtypes.double) + indices = delta.cumsum(0).cast(dtypes.int64) + + # input = torch.randn(indices.abs().max() + 1) + input = Tensor.randn(indices.abs().max().item() + 1) + # values = torch.randn(indices.size(0)) + values = Tensor.randn(indices.shape(0)) + output = input.index_put((indices,), values, accumulate=True) + + input_list = input.numpy().tolist() + indices_list = indices.numpy().tolist() + values_list = values.numpy().tolist() + for i, v in zip(indices_list, values_list): + input_list[i] += v + + numpy_testing_assert_equal_helper(output, input_list) + ''' + + def test_index_ind_dtype(self): + x = Tensor.randn(4, 4) + # ind_long = torch.randint(4, (4,), dtype=torch.long) + # TODO should we spend an extra line to allow for randint other dtypes? + # copied from randint + ind_long = (Tensor.rand((4,),)*(4-0)+0).cast(dtypes.int64) + # ind_int = ind_long.int() + ind_int = (ind_long).cast(dtypes.int32) + ref = x[ind_long, ind_long] + res = x[ind_int, ind_int] + numpy_testing_assert_equal_helper(ref, res) + ref = x[ind_long, :] + res = x[ind_int, :] + numpy_testing_assert_equal_helper(ref, res) + ref = x[:, ind_long] + res = x[:, ind_int] + numpy_testing_assert_equal_helper(ref, res) + # no repeating indices for index_put + # TODO setitem + ''' + src = Tensor.randn(4) + ind_long = Tensor.arange(4, dtype=dtypes.int64) + ind_int = ind_long.cast(dtypes.int32) + for accum in (True, False): + inp_ref = clone(x) + inp_res = clone(x) + torch.index_put_(inp_ref, (ind_long, ind_long), src, accum) + torch.index_put_(inp_res, (ind_int, ind_int), src, accum) + numpy_testing_assert_equal_helper(inp_ref, inp_res) + ''' + + # TODO setitem + ''' + def test_index_put_accumulate_empty(self): + # Regression test for https://github.com/pytorch/pytorch/issues/94667 + input = Tensor.rand([], dtype=dtypes.float32) + with self.assertRaises(RuntimeError): + input.index_put([], np.array([1.0]), True) + ''' + + # TODO bool indexing + ''' + def test_multiple_byte_mask(self): + v = Tensor.randn(5, 7, 3) + # note: these broadcast together and are transposed to the first dim + mask1 = Tensor([1, 0, 1, 1, 0], dtype=dtypes.uint8) + mask2 = Tensor([1, 1, 1], dtype=dtypes.uint8) + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + numpy_testing_assert_equal_helper(v[mask1, :, mask2].shape, (3, 7)) + numpy_testing_assert_equal_helper(len(w), 2) + ''' + + # TODO bool indexing + ''' + def test_byte_mask2d(self): + v = Tensor.randn(5, 7, 3) + c = Tensor.randn(5, 7) + num_ones = (c > 0).sum() + r = v[c > 0] + numpy_testing_assert_equal_helper(r.shape, (num_ones, 3)) + ''' + + # TODO setindex + # TODO bool indexing + ''' + def test_jit_indexing(self): + def fn1(x): + x[x < 50] = 1.0 + return x + + def fn2(x): + x[0:50] = 1.0 + return x + + scripted_fn1 = TinyJit(fn1) + scripted_fn2 = TinyJit(fn2) + data = Tensor.arange(100, dtype=dtypes.float) + out = scripted_fn1(clone(data)) + ref = Tensor(np.concatenate((np.ones(50), np.arange(50, 100))), dtype=dtypes.float) + numpy_testing_assert_equal_helper(out, ref) + out = scripted_fn2(clone(data)) + numpy_testing_assert_equal_helper(out, ref) + ''' + + def test_int_indices(self): + v = Tensor.randn(5, 7, 3) + numpy_testing_assert_equal_helper(v[[0, 4, 2]].shape, (3, 7, 3)) + numpy_testing_assert_equal_helper(v[:, [0, 4, 2]].shape, (5, 3, 3)) + numpy_testing_assert_equal_helper(v[:, [[0, 1], [4, 3]]].shape, (5, 2, 2, 3)) + + # TODO setindex + ''' + def test_index_put_src_datatype(self, dtype): + src = Tensor.ones(3, 2, 4, dtype=dtype) + vals = Tensor.ones(3, 2, 4, dtype=dtype) + indices = (np.array([0, 2, 1]),) + res = src.index_put_(indices, vals, accumulate=True) + numpy_testing_assert_equal_helper(res.shape, src.shape) + ''' + + def test_index_src_datatype(self): + src = Tensor.ones(3, 2, 4) + # test index + res = src[[0, 2, 1], :, :] + numpy_testing_assert_equal_helper(res.shape, src.shape) + # test index_put, no accum + # TODO setindex + ''' + src[[0, 2, 1], :, :] = res + numpy_testing_assert_equal_helper(res.shape, src.shape) + ''' + + def test_int_indices2d(self): + # From the NumPy indexing example + x = Tensor.arange(0, 12).reshape(4, 3) + rows = Tensor([[0, 0], [3, 3]]) + columns = Tensor([[0, 2], [0, 2]]) + numpy_testing_assert_equal_helper(x[rows, columns].numpy().tolist(), [[0, 2], [9, 11]]) + + def test_int_indices_broadcast(self): + # From the NumPy indexing example + x = Tensor.arange(0, 12).reshape(4, 3) + rows = Tensor([0, 3]) + columns = Tensor([0, 2]) + result = x[rows[:, None], columns] + numpy_testing_assert_equal_helper(result.numpy().tolist(), [[0, 2], [9, 11]]) + + # TODO setitem + # TODO empty Tensor fancy index + ''' + def test_empty_index(self): + x = Tensor.arange(0, 12).reshape(4, 3) + idx = Tensor([], dtype=dtypes.int64) + numpy_testing_assert_equal_helper(x[idx].numel(), 0) + + # empty assignment should have no effect but not throw an exception + y = clone(x) + y[idx] = -1 + numpy_testing_assert_equal_helper(x, y) + + mask = Tensor.zeros(4, 3).cast(dtypes.bool) + y[mask] = -1 + numpy_testing_assert_equal_helper(x, y) + ''' + + # TODO empty Tensor fancy index + ''' + def test_empty_ndim_index(self): + x = Tensor.randn(5) + numpy_testing_assert_equal_helper(Tensor.empty(0, 2), x[Tensor.empty(0, 2, dtype=dtypes.int64)]) + + x = Tensor.randn(2, 3, 4, 5) + numpy_testing_assert_equal_helper(Tensor.empty(2, 0, 6, 4, 5), + x[:, Tensor.empty(0, 6, dtype=dtypes.int64)]) + + x = Tensor.empty(10, 0) + numpy_testing_assert_equal_helper(x[[1, 2]].shape, (2, 0)) + numpy_testing_assert_equal_helper(x[[], []].shape, (0,)) + with self.assertRaisesRegex(IndexError, 'for dimension with size 0'): + x[:, [0, 1]] + ''' + + # TODO empty Tensor fancy index + ''' + def test_empty_ndim_index_bool(self): + x = Tensor.randn(5) + self.assertRaises(IndexError, lambda: x[Tensor.empty(0, 2, dtype=dtypes.uint8)]) + ''' + + def test_empty_slice(self): + x = Tensor.randn(2, 3, 4, 5) + y = x[:, :, :, 1] + z = y[:, 1:1, :] + numpy_testing_assert_equal_helper((2, 0, 4), z.shape) + # this isn't technically necessary, but matches NumPy stride calculations. + # TODO not too sure about this + # numpy_testing_assert_equal_helper((60, 20, 5), z.lazydata.st.real_strides()) + self.assertTrue(z.lazydata.st.contiguous) + + # TODO bool indexing + # TODO data_ptr() + ''' + def test_index_getitem_copy_bools_slices(self): + true = Tensor(1, dtype=dtypes.uint8) + false = Tensor(0, dtype=dtypes.uint8) + + tensors = [Tensor.randn(2, 3), Tensor(3.)] + + for a in tensors: + self.assertNotEqual(a.data_ptr(), a[True].data_ptr()) + numpy_testing_assert_equal_helper(Tensor.empty(0, *a.shape), a[False]) + self.assertNotEqual(a.data_ptr(), a[true].data_ptr()) + numpy_testing_assert_equal_helper(Tensor.empty(0, *a.shape), a[false]) + numpy_testing_assert_equal_helper(a.data_ptr(), a[None].data_ptr()) + numpy_testing_assert_equal_helper(a.data_ptr(), a[...].data_ptr()) + ''' + + # TODO setitem + # TODO bool indexing + ''' + def test_index_setitem_bools_slices(self): + true = Tensor(1, dtype=dtypes.uint8) + false = Tensor(0, dtype=dtypes.uint8) + + tensors = [Tensor.randn(2, 3), Tensor(3)] + + for a in tensors: + # prefix with a 1,1, to ensure we are compatible with numpy which cuts off prefix 1s + # (some of these ops already prefix a 1 to the size) + neg_ones = Tensor.ones_like(a) * -1 + neg_ones_expanded = neg_ones.unsqueeze(0).unsqueeze(0) + a[True] = neg_ones_expanded + numpy_testing_assert_equal_helper(a, neg_ones) + a[False] = 5 + numpy_testing_assert_equal_helper(a, neg_ones) + a[true] = neg_ones_expanded * 2 + numpy_testing_assert_equal_helper(a, neg_ones * 2) + a[false] = 5 + numpy_testing_assert_equal_helper(a, neg_ones * 2) + a[None] = neg_ones_expanded * 3 + numpy_testing_assert_equal_helper(a, neg_ones * 3) + a[...] = neg_ones_expanded * 4 + numpy_testing_assert_equal_helper(a, neg_ones * 4) + if a.dim() == 0: + with self.assertRaises(IndexError): + a[:] = neg_ones_expanded * 5 + ''' + + # TODO bool indexing + ''' + def test_index_scalar_with_bool_mask(self): + a = Tensor(1) + uintMask = Tensor(True, dtype=dtypes.uint8) + boolMask = Tensor(True, dtype=dtypes.bool) + numpy_testing_assert_equal_helper(a[uintMask], a[boolMask]) + numpy_testing_assert_equal_helper(a[uintMask].dtype, a[boolMask].dtype) + + a = Tensor(True, dtype=dtypes.bool) + numpy_testing_assert_equal_helper(a[uintMask], a[boolMask]) + numpy_testing_assert_equal_helper(a[uintMask].dtype, a[boolMask].dtype) + ''' + + # TODO setitem + ''' + def test_setitem_expansion_error(self): + true = Tensor(True) + a = Tensor.randn(2, 3) + # check prefix with non-1s doesn't work + # a_expanded = a.expand(torch.Size([5, 1]) + a.size()) + a_expanded = a.expand((5, 1) + a.shape) + # NumPy: ValueError + with self.assertRaises(RuntimeError): + a[True] = a_expanded + with self.assertRaises(RuntimeError): + a[true] = a_expanded + ''' + + def test_getitem_scalars(self): + zero = Tensor(0, dtype=dtypes.int64) + one = Tensor(1, dtype=dtypes.int64) + + # non-scalar indexed with scalars + a = Tensor.randn(2, 3) + numpy_testing_assert_equal_helper(a[0], a[zero]) + numpy_testing_assert_equal_helper(a[0][1], a[zero][one]) + numpy_testing_assert_equal_helper(a[0, 1], a[zero, one]) + numpy_testing_assert_equal_helper(a[0, one], a[zero, 1]) + + # indexing by a scalar should slice (not copy) + # TODO data ptr + ''' + numpy_testing_assert_equal_helper(a[0, 1].data_ptr(), a[zero, one].data_ptr()) + numpy_testing_assert_equal_helper(a[1].data_ptr(), a[one.cast(dtypes.int32)].data_ptr()) + numpy_testing_assert_equal_helper(a[1].data_ptr(), a[one.cast(dtypes.int16)].data_ptr()) + ''' + + # scalar indexed with scalar + r = Tensor.randn() + with self.assertRaises(IndexError): + r[:] + with self.assertRaises(IndexError): + r[zero] + numpy_testing_assert_equal_helper(r, r[...]) + + # TODO setitem + ''' + def test_setitem_scalars(self): + zero = Tensor(0, dtype=dtypes.int64) + + # non-scalar indexed with scalars + a = Tensor.randn(2, 3) + a_set_with_number = clone(a) + a_set_with_scalar = clone(a) + b = Tensor.randn(3) + + a_set_with_number[0] = b + a_set_with_scalar[zero] = b + numpy_testing_assert_equal_helper(a_set_with_number, a_set_with_scalar) + a[1, zero] = 7.7 + numpy_testing_assert_equal_helper(7.7, a[1, 0]) + + # scalar indexed with scalars + r = Tensor.randn() + with self.assertRaises(IndexError): + r[:] = 8.8 + with self.assertRaises(IndexError): + r[zero] = 8.8 + r[...] = 9.9 + numpy_testing_assert_equal_helper(9.9, r) + ''' + + def test_basic_advanced_combined(self): + # From the NumPy indexing example + x = Tensor.arange(0, 12).reshape(4, 3) + numpy_testing_assert_equal_helper(x[1:2, 1:3], x[1:2, [1, 2]]) + numpy_testing_assert_equal_helper(x[1:2, 1:3].numpy().tolist(), [[4, 5]]) + + # Check that it is a copy + unmodified = clone(x) + # x[1:2, [1, 2]].zero_() + x[1:2, [1, 2]].zeros_like() + numpy_testing_assert_equal_helper(x, unmodified) + + # But assignment should modify the original + # TODO setitem + ''' + unmodified = clone(x) + x[1:2, [1, 2]] = 0 + self.assertNotEqual(x, unmodified) + ''' + + # TODO setitem + ''' + def test_int_assignment(self): + x = Tensor.arange(0, 4).reshape(2, 2) + x[1] = 5 + numpy_testing_assert_equal_helper(x.numpy().tolist(), [[0, 1], [5, 5]]) + + x = Tensor.arange(0, 4).reshape(2, 2) + x[1] = Tensor.arange(5, 7) + numpy_testing_assert_equal_helper(x.numpy().tolist(), [[0, 1], [5, 6]]) + ''' + + # TODO setitem + ''' + def test_byte_tensor_assignment(self): + x = Tensor.arange(0., 16).reshape(4, 4) + b = Tensor([True, False, True, False], dtype=dtypes.uint8) + value = Tensor([3., 4., 5., 6.]) + + with warnings.catch_warnings(record=True) as w: + x[b] = value + numpy_testing_assert_equal_helper(len(w), 1) + + numpy_testing_assert_equal_helper(x[0], value) + numpy_testing_assert_equal_helper(x[1], Tensor.arange(4., 8)) + numpy_testing_assert_equal_helper(x[2], value) + numpy_testing_assert_equal_helper(x[3], Tensor.arange(12., 16)) + ''' + + # TODO tensor unpacking + ''' + def test_variable_slicing(self): + x = Tensor.arange(0, 16).reshape(4, 4) + indices = Tensor([0, 1], dtype=dtypes.int32) + i, j = indices + numpy_testing_assert_equal_helper(x[i:j], x[0:1]) + ''' + + def test_ellipsis_tensor(self): + x = Tensor.arange(0, 9).reshape(3, 3) + idx = Tensor([0, 2]) + numpy_testing_assert_equal_helper(x[..., idx].numpy().tolist(), [[0, 2], + [3, 5], + [6, 8]]) + numpy_testing_assert_equal_helper(x[idx, ...].numpy().tolist(), [[0, 1, 2], + [6, 7, 8]]) + + # TODO unravel_index # def test_unravel_index_errors(self): - # with self.assertRaisesRegex(TypeError, r"expected 'indices' to be integer"): - # torch.unravel_index( - # np.array(0.5), - # (2, 2)) + # with self.assertRaisesRegex(TypeError, r"expected 'indices' to be integer"): + # unravel_index( + # Tensor(0.5), + # (2, 2)) - # with self.assertRaisesRegex(TypeError, r"expected 'indices' to be integer"): - # torch.unravel_index( - # np.array([]), - # (10, 3, 5)) + # with self.assertRaisesRegex(TypeError, r"expected 'indices' to be integer"): + # unravel_index( + # Tensor([]), + # (10, 3, 5)) - # with self.assertRaisesRegex(TypeError, r"expected 'shape' to be int or sequence"): - # torch.unravel_index( - # np.array([1], dtype=torch.int64), - # np.array([1, 2, 3])) + # with self.assertRaisesRegex(TypeError, r"expected 'shape' to be int or sequence"): + # unravel_index( + # Tensor([1], dtype=dtypes.int64), + # Tensor([1, 2, 3])) - # with self.assertRaisesRegex(TypeError, r"expected 'shape' sequence to only contain ints"): - # torch.unravel_index( - # np.array([1], dtype=torch.int64), - # (1, 2, 2.0)) + # with self.assertRaisesRegex(TypeError, r"expected 'shape' sequence to only contain ints"): + # unravel_index( + # Tensor([1], dtype=dtypes.int64), + # (1, 2, 2.0)) - # with self.assertRaisesRegex(ValueError, r"'shape' cannot have negative values, but got \(2, -3\)"): - # torch.unravel_index( - # np.array(0), - # (2, -3)) + # with self.assertRaisesRegex(ValueError, r"'shape' cannot have negative values, but got \(2, -3\)"): + # unravel_index( + # Tensor(0), + # (2, -3)) - # def test_invalid_index(self): - # x = torch.arange(0, 16).view(4, 4) - # self.assertRaisesRegex(TypeError, 'slice indices', lambda: x["0":"1"]) + def test_invalid_index(self): + x = Tensor.arange(0, 16).reshape(4, 4) + self.assertRaisesRegex(TypeError, 'slice indices', lambda: x["0":"1"]) - # def test_out_of_bound_index(self): - # x = torch.arange(0, 100).view(2, 5, 10) - # self.assertRaisesRegex(IndexError, 'index 5 is out of bounds for dimension 1 with size 5', lambda: x[0, 5]) - # self.assertRaisesRegex(IndexError, 'index 4 is out of bounds for dimension 0 with size 2', lambda: x[4, 5]) - # self.assertRaisesRegex(IndexError, 'index 15 is out of bounds for dimension 2 with size 10', - # lambda: x[0, 1, 15]) - # self.assertRaisesRegex(IndexError, 'index 12 is out of bounds for dimension 2 with size 10', - # lambda: x[:, :, 12]) + def test_out_of_bound_index(self): + x = Tensor.arange(0, 100).reshape(2, 5, 10) + self.assertRaisesRegex(IndexError, 'index 5 is out of bounds for dimension 1 with size 5', lambda: x[0, 5]) + self.assertRaisesRegex(IndexError, 'index 4 is out of bounds for dimension 0 with size 2', lambda: x[4, 5]) + self.assertRaisesRegex(IndexError, 'index 15 is out of bounds for dimension 2 with size 10', lambda: x[0, 1, 15]) + self.assertRaisesRegex(IndexError, 'index 12 is out of bounds for dimension 2 with size 10', lambda: x[:, :, 12]) - # def test_zero_dim_index(self): - # x = np.array(10) - # numpy_testing_assert_equal_helper(x, x.item()) + # TODO .item() bug + ''' + def test_zero_dim_index(self): + x = Tensor(10) + numpy_testing_assert_equal_helper(x, x.item()) - # def runner(): - # print(x[0]) - # return x[0] + def runner(): + print(x[0]) + return x[0] - # self.assertRaisesRegex(IndexError, 'invalid index', runner) + self.assertRaisesRegex(IndexError, 'invalid index', runner) + ''' - # def test_invalid_device(self): - # idx = np.array([0, 1]) - # b = torch.zeros(5) - # c = np.array([1., 2.], device="cpu") + # TODO not too sure + ''' + def test_invalid_device(self): + idx = Tensor([0, 1]) + b = Tensor.zeros(5) + c = Tensor([1., 2.], device="CPU") - # for accumulate in [True, False]: - # self.assertRaises(RuntimeError, lambda: torch.index_put_(b, (idx,), c, accumulate=accumulate)) + for accumulate in [True, False]: + self.assertRaises(RuntimeError, lambda: torch.index_put_(b, (idx,), c, accumulate=accumulate)) + ''' - # def test_cpu_indices(self): - # idx = np.array([0, 1]) - # b = torch.zeros(2) - # x = torch.ones(10) - # x[idx] = b # index_put_ - # ref = torch.ones(10) - # ref[:2] = 0 - # numpy_testing_assert_equal_helper(x, ref) - # out = x[idx] # index - # numpy_testing_assert_equal_helper(out, torch.zeros(2)) + # TODO setitem + ''' + def test_cpu_indices(self): + idx = Tensor([0, 1]) + b = Tensor.zeros(2) + x = Tensor.ones(10) + x[idx] = b # index_put_ + ref = Tensor.ones(10) + ref[:2] = 0 + numpy_testing_assert_equal_helper(x, ref) + out = x[idx] # index + numpy_testing_assert_equal_helper(out, Tensor.zeros(2)) + ''' def test_take_along_dim(self): ''' @@ -1287,144 +1440,155 @@ class TestIndexing(unittest.TestCase): expected = np.take_along_axis(t_np, indices_np, axis=dim) numpy_testing_assert_equal_helper(actual, expected) - # for shape in [(3, 2), (2, 3, 5), (2, 4, 0), (2, 3, 1, 4)]: - # for noncontiguous in [True, False]: - # t = make_tensor(shape, dtype=dtype, noncontiguous=noncontiguous) - # for dim in list(range(t.ndim)) + [None]: - # if dim is None: - # indices = torch.argsort(t.view(-1)) - # else: - # indices = torch.argsort(t, dim=dim) + # TODO argsort + ''' + for shape in [(3, 2), (2, 3, 5), (2, 4, 0), (2, 3, 1, 4)]: + for noncontiguous in [True, False]: + for dtype in (dtypes.float32, dtypes.int64): + t = make_tensor(shape, dtype=dtype, noncontiguous=noncontiguous) + for dim in list(range(t.ndim)) + [None]: + if dim is None: + indices = argsort(t.reshape(-1)) + else: + indices = argsort(t, dim=dim) - # _test_against_numpy(t, indices, dim) + _test_against_numpy(t, indices, dim) + ''' - # # test broadcasting - # t = torch.ones((3, 4, 1)) - # indices = torch.ones((1, 2, 5), dtype=torch.long) + # test broadcasting + t = Tensor.ones((3, 4, 1)) + indices = Tensor.ones((1, 2, 5), dtype=dtypes.int64) - # _test_against_numpy(t, indices, 1) + _test_against_numpy(t, indices, 1) - # # test empty indices - # t = torch.ones((3, 4, 5)) - # indices = torch.ones((3, 0, 5), dtype=torch.long) + # test empty indices + t = Tensor.ones((3, 4, 5)) + indices = Tensor.ones((3, 0, 5), dtype=dtypes.int64) - # _test_against_numpy(t, indices, 1) + _test_against_numpy(t, indices, 1) - # def test_take_along_dim_invalid(self, dtype): - # shape = (2, 3, 1, 4) - # dim = 0 - # t = make_tensor(shape, dtype=dtype) - # indices = torch.argsort(t, dim=dim) + # TODO argsort + ''' + def test_take_along_dim_invalid(self): + for dtype in (dtypes.int64, dtypes.float32): + shape = (2, 3, 1, 4) + dim = 0 + t = make_tensor(shape, dtype=dtype) + indices = argsort(t, dim=dim) - # # dim of `t` and `indices` does not match - # with self.assertRaisesRegex(RuntimeError, - # "input and indices should have the same number of dimensions"): - # torch.take_along_dim(t, indices[0], dim=0) + # dim of `t` and `indices` does not match + with self.assertRaisesRegex(RuntimeError, "input and indices should have the same number of dimensions"): + t.gather(indices[0], dim=0) - # # invalid `indices` dtype - # with self.assertRaisesRegex(RuntimeError, r"dtype of indices should be Long"): - # torch.take_along_dim(t, indices.to(torch.bool), dim=0) + # invalid `indices` dtype + with self.assertRaisesRegex(RuntimeError, r"dtype of indices should be Long"): + t.gather(indices.cast(dtypes.bool), dim=0) - # with self.assertRaisesRegex(RuntimeError, r"dtype of indices should be Long"): - # torch.take_along_dim(t, indices.to(torch.float), dim=0) + with self.assertRaisesRegex(RuntimeError, r"dtype of indices should be Long"): + t.gather(indices.cast(dtypes.float32), dim=0) - # with self.assertRaisesRegex(RuntimeError, r"dtype of indices should be Long"): - # torch.take_along_dim(t, indices.to(torch.int32), dim=0) + with self.assertRaisesRegex(RuntimeError, r"dtype of indices should be Long"): + t.gather(indices.cast(dtypes.int32), dim=0) - # # invalid axis - # with self.assertRaisesRegex(IndexError, "Dimension out of range"): - # torch.take_along_dim(t, indices, dim=-7) + # invalid axis + with self.assertRaisesRegex(IndexError, "Dimension out of range"): + t.gather(indices, dim=-7) - # with self.assertRaisesRegex(IndexError, "Dimension out of range"): - # torch.take_along_dim(t, indices, dim=7) + with self.assertRaisesRegex(IndexError, "Dimension out of range"): + t.gather(t, indices, dim=7) + ''' - # def test_gather_take_along_dim_cross_device(self, dtype): - # shape = (2, 3, 1, 4) - # dim = 0 - # t = make_tensor(shape, dtype=dtype) - # indices = torch.argsort(t, dim=dim) + # TODO Exception for different devices + # TODO argsort + ''' + def test_gather_take_along_dim_cross_device(self, dtype=dtypes.float32): + shape = (2, 3, 1, 4) + dim = 0 + t = make_tensor(shape, dtype=dtype) + indices = argsort(t, dim=dim) - # with self.assertRaisesRegex(RuntimeError, "Expected all tensors to be on the same device"): - # torch.gather(t, 0, indices.cpu()) + with self.assertRaisesRegex(RuntimeError, "Expected all tensors to be on the same device"): + torch.gather(t, 0, indices.cpu()) - # with self.assertRaisesRegex(RuntimeError, - # r"Expected tensor to have .* but got tensor with .* torch.take_along_dim()"): - # torch.take_along_dim(t, indices.cpu(), dim=0) + with self.assertRaisesRegex(RuntimeError, r"Expected tensor to have .* but got tensor with .* torch.take_along_dim()"): + torch.take_along_dim(t, indices.cpu(), dim=0) - # with self.assertRaisesRegex(RuntimeError, "Expected all tensors to be on the same device"): - # torch.gather(t.cpu(), 0, indices) + with self.assertRaisesRegex(RuntimeError, "Expected all tensors to be on the same device"): + torch.gather(t.cpu(), 0, indices) - # with self.assertRaisesRegex(RuntimeError, - # r"Expected tensor to have .* but got tensor with .* torch.take_along_dim()"): - # torch.take_along_dim(t.cpu(), indices, dim=0) + with self.assertRaisesRegex(RuntimeError, r"Expected tensor to have .* but got tensor with .* torch.take_along_dim()"): + torch.take_along_dim(t.cpu(), indices, dim=0) + ''' - # def test_cuda_broadcast_index_use_deterministic_algorithms(self): - # with DeterministicGuard(True): - # idx1 = np.array([0]) - # idx2 = np.array([2, 6]) - # idx3 = np.array([1, 5, 7]) + # TODO another torch specific test... + ''' + def test_cuda_broadcast_index_use_deterministic_algorithms(self): + with DeterministicGuard(True): + idx1 = np.array([0]) + idx2 = np.array([2, 6]) + idx3 = np.array([1, 5, 7]) - # tensor_a = torch.rand(13, 11, 12, 13, 12).cpu() - # tensor_b = tensor_a.to(device=device) - # tensor_a[idx1] = 1.0 - # tensor_a[idx1, :, idx2, idx2, :] = 2.0 - # tensor_a[:, idx1, idx3, :, idx3] = 3.0 - # tensor_b[idx1] = 1.0 - # tensor_b[idx1, :, idx2, idx2, :] = 2.0 - # tensor_b[:, idx1, idx3, :, idx3] = 3.0 - # numpy_testing_assert_equal_helper(tensor_a, tensor_b.cpu()) + tensor_a = torch.rand(13, 11, 12, 13, 12).cpu() + tensor_b = tensor_a.to(device=device) + tensor_a[idx1] = 1.0 + tensor_a[idx1, :, idx2, idx2, :] = 2.0 + tensor_a[:, idx1, idx3, :, idx3] = 3.0 + tensor_b[idx1] = 1.0 + tensor_b[idx1, :, idx2, idx2, :] = 2.0 + tensor_b[:, idx1, idx3, :, idx3] = 3.0 + numpy_testing_assert_equal_helper(tensor_a, tensor_b.cpu()) - # tensor_a = torch.rand(10, 11).cpu() - # tensor_b = tensor_a.to(device=device) - # tensor_a[idx3] = 1.0 - # tensor_a[idx2, :] = 2.0 - # tensor_a[:, idx2] = 3.0 - # tensor_a[:, idx1] = 4.0 - # tensor_b[idx3] = 1.0 - # tensor_b[idx2, :] = 2.0 - # tensor_b[:, idx2] = 3.0 - # tensor_b[:, idx1] = 4.0 - # numpy_testing_assert_equal_helper(tensor_a, tensor_b.cpu()) + tensor_a = torch.rand(10, 11).cpu() + tensor_b = tensor_a.to(device=device) + tensor_a[idx3] = 1.0 + tensor_a[idx2, :] = 2.0 + tensor_a[:, idx2] = 3.0 + tensor_a[:, idx1] = 4.0 + tensor_b[idx3] = 1.0 + tensor_b[idx2, :] = 2.0 + tensor_b[:, idx2] = 3.0 + tensor_b[:, idx1] = 4.0 + numpy_testing_assert_equal_helper(tensor_a, tensor_b.cpu()) - # tensor_a = torch.rand(10, 10).cpu() - # tensor_b = tensor_a.to(device=device) - # tensor_a[[8]] = 1.0 - # tensor_b[[8]] = 1.0 - # numpy_testing_assert_equal_helper(tensor_a, tensor_b.cpu()) - - # tensor_a = torch.rand(10).cpu() - # tensor_b = tensor_a.to(device=device) - # tensor_a[6] = 1.0 - # tensor_b[6] = 1.0 - # numpy_testing_assert_equal_helper(tensor_a, tensor_b.cpu()) + tensor_a = torch.rand(10, 10).cpu() + tensor_b = tensor_a.to(device=device) + tensor_a[[8]] = 1.0 + tensor_b[[8]] = 1.0 + numpy_testing_assert_equal_helper(tensor_a, tensor_b.cpu()) + tensor_a = torch.rand(10).cpu() + tensor_b = tensor_a.to(device=device) + tensor_a[6] = 1.0 + tensor_b[6] = 1.0 + numpy_testing_assert_equal_helper(tensor_a, tensor_b.cpu()) + ''' class TestNumpy(unittest.TestCase): - # def test_index_no_floats(self): - # a = Tensor([[[5.]]]) + def test_index_no_floats(self): + a = Tensor([[[5.]]]) - # self.assertRaises(IndexError, lambda: a[0.0]) - # self.assertRaises(IndexError, lambda: a[0, 0.0]) - # self.assertRaises(IndexError, lambda: a[0.0, 0]) - # self.assertRaises(IndexError, lambda: a[0.0, :]) - # self.assertRaises(IndexError, lambda: a[:, 0.0]) - # self.assertRaises(IndexError, lambda: a[:, 0.0, :]) - # self.assertRaises(IndexError, lambda: a[0.0, :, :]) - # self.assertRaises(IndexError, lambda: a[0, 0, 0.0]) - # self.assertRaises(IndexError, lambda: a[0.0, 0, 0]) - # self.assertRaises(IndexError, lambda: a[0, 0.0, 0]) - # self.assertRaises(IndexError, lambda: a[-1.4]) - # self.assertRaises(IndexError, lambda: a[0, -1.4]) - # self.assertRaises(IndexError, lambda: a[-1.4, 0]) - # self.assertRaises(IndexError, lambda: a[-1.4, :]) - # self.assertRaises(IndexError, lambda: a[:, -1.4]) - # self.assertRaises(IndexError, lambda: a[:, -1.4, :]) - # self.assertRaises(IndexError, lambda: a[-1.4, :, :]) - # self.assertRaises(IndexError, lambda: a[0, 0, -1.4]) - # self.assertRaises(IndexError, lambda: a[-1.4, 0, 0]) - # self.assertRaises(IndexError, lambda: a[0, -1.4, 0]) - # # self.assertRaises(IndexError, lambda: a[0.0:, 0.0]) - # # self.assertRaises(IndexError, lambda: a[0.0:, 0.0,:]) + self.assertRaises(IndexError, lambda: a[0.0]) + self.assertRaises(IndexError, lambda: a[0, 0.0]) + self.assertRaises(IndexError, lambda: a[0.0, 0]) + self.assertRaises(IndexError, lambda: a[0.0, :]) + self.assertRaises(IndexError, lambda: a[:, 0.0]) + self.assertRaises(IndexError, lambda: a[:, 0.0, :]) + self.assertRaises(IndexError, lambda: a[0.0, :, :]) + self.assertRaises(IndexError, lambda: a[0, 0, 0.0]) + self.assertRaises(IndexError, lambda: a[0.0, 0, 0]) + self.assertRaises(IndexError, lambda: a[0, 0.0, 0]) + self.assertRaises(IndexError, lambda: a[-1.4]) + self.assertRaises(IndexError, lambda: a[0, -1.4]) + self.assertRaises(IndexError, lambda: a[-1.4, 0]) + self.assertRaises(IndexError, lambda: a[-1.4, :]) + self.assertRaises(IndexError, lambda: a[:, -1.4]) + self.assertRaises(IndexError, lambda: a[:, -1.4, :]) + self.assertRaises(IndexError, lambda: a[-1.4, :, :]) + self.assertRaises(IndexError, lambda: a[0, 0, -1.4]) + self.assertRaises(IndexError, lambda: a[-1.4, 0, 0]) + self.assertRaises(IndexError, lambda: a[0, -1.4, 0]) + self.assertRaises(IndexError, lambda: a[0.0:, 0.0]) + self.assertRaises(IndexError, lambda: a[0.0:, 0.0,:]) def test_none_index(self): # `None` index adds newaxis @@ -1435,43 +1599,52 @@ class TestNumpy(unittest.TestCase): # Empty tuple index creates a view a = Tensor([1, 2, 3]) numpy_testing_assert_equal_helper(a[()], a) - # # TODO: what's our equivalent test? just is? - # numpy_testing_assert_equal_helper(a[()].data_ptr(), a.data_ptr()) + # TODO data_ptr + ''' + numpy_testing_assert_equal_helper(a[()].data_ptr(), a.data_ptr()) + ''' - # def test_empty_fancy_index(self): - # # Empty list index creates an empty array - # a = Tensor([1, 2, 3]) - # numpy_testing_assert_equal_helper(a[[]], np.array([])) + # TODO empty fancy index + ''' + def test_empty_fancy_index(self): + # Empty list index creates an empty array + a = Tensor([1, 2, 3]) + numpy_testing_assert_equal_helper(a[[]], np.array([])) - # b = Tensor([]).long() - # numpy_testing_assert_equal_helper(a[[]], np.array([])) + b = Tensor([]).cast(dtypes.int64) + numpy_testing_assert_equal_helper(a[[]], np.array([])) - # b = Tensor([]).float() - # self.assertRaises(IndexError, lambda: a[b]) + # TODO fancy index dtype error + b = Tensor([]).float() + self.assertRaises(IndexError, lambda: a[b]) + ''' -# def test_ellipsis_index(self): -# a = tensor([[1, 2, 3], -# [4, 5, 6], -# [7, 8, 9]]) -# self.assertIsNot(a[...], a) -# numpy_testing_assert_equal_helper(a[...], a) -# # `a[...]` was `a` in numpy <1.9. -# numpy_testing_assert_equal_helper(a[...].data_ptr(), a.data_ptr()) + def test_ellipsis_index(self): + a = Tensor([[1, 2, 3], + [4, 5, 6], + [7, 8, 9]]) + self.assertIsNot(a[...], a) + numpy_testing_assert_equal_helper(a[...], a) + # `a[...]` was `a` in numpy <1.9. + # TODO data_ptr + ''' + numpy_testing_assert_equal_helper(a[...].data_ptr(), a.data_ptr()) + ''' -# # Slicing with ellipsis can skip an -# # arbitrary number of dimensions -# numpy_testing_assert_equal_helper(a[0, ...], a[0]) -# numpy_testing_assert_equal_helper(a[0, ...], a[0, :]) -# numpy_testing_assert_equal_helper(a[..., 0], a[:, 0]) + # Slicing with ellipsis can skip an + # arbitrary number of dimensions + numpy_testing_assert_equal_helper(a[0, ...], a[0]) + numpy_testing_assert_equal_helper(a[0, ...], a[0, :]) + numpy_testing_assert_equal_helper(a[..., 0], a[:, 0]) -# # In NumPy, slicing with ellipsis results in a 0-dim array. In PyTorch -# # we don't have separate 0-dim arrays and scalars. -# numpy_testing_assert_equal_helper(a[0, ..., 1], np.array(2)) + # In NumPy, slicing with ellipsis results in a 0-dim array. In PyTorch + # we don't have separate 0-dim arrays and scalars. + numpy_testing_assert_equal_helper(a[0, ..., 1], np.array(2)) -# # Assignment with `(Ellipsis,)` on 0-d arrays -# b = np.array(1) -# b[(Ellipsis,)] = 2 -# numpy_testing_assert_equal_helper(b, 2) + # Assignment with `(Ellipsis,)` on 0-d arrays + b = np.array(1) + b[(Ellipsis,)] = 2 + numpy_testing_assert_equal_helper(b, 2) def test_single_int_index(self): # Single integer index selects one row @@ -1485,154 +1658,200 @@ class TestNumpy(unittest.TestCase): self.assertRaises(IndexError, a.__getitem__, 1 << 30) self.assertRaises(IndexError, a.__getitem__, 1 << 64) - # def test_single_bool_index(self): - # # Single boolean index - # a = Tensor([[1, 2, 3], - # [4, 5, 6], - # [7, 8, 9]]) + # TODO bool indexing + ''' + def test_single_bool_index(self): + # Single boolean index + a = Tensor([[1, 2, 3], + [4, 5, 6], + [7, 8, 9]]) - # numpy_testing_assert_equal_helper(a[True], a[None]) - # numpy_testing_assert_equal_helper(a[False], a[None][0:0]) + numpy_testing_assert_equal_helper(a[True], a[None]) + numpy_testing_assert_equal_helper(a[False], a[None][0:0]) + ''' -# def test_boolean_shape_mismatch(self): -# arr = torch.ones((5, 4, 3)) + # TODO bool indexing + ''' + def test_boolean_shape_mismatch(self): + arr = Tensor.ones((5, 4, 3)) -# index = tensor([True]) -# self.assertRaisesRegex(IndexError, 'mask', lambda: arr[index]) + index = Tensor([True]) + self.assertRaisesRegex(IndexError, 'mask', lambda: arr[index]) -# index = tensor([False] * 6) -# self.assertRaisesRegex(IndexError, 'mask', lambda: arr[index]) + index = Tensor([False] * 6) + self.assertRaisesRegex(IndexError, 'mask', lambda: arr[index]) -# index = torch.ByteTensor(4, 4).to(device).zero_() -# self.assertRaisesRegex(IndexError, 'mask', lambda: arr[index]) -# self.assertRaisesRegex(IndexError, 'mask', lambda: arr[(slice(None), index)]) + # index = torch.ByteTensor(4, 4).to(device).zero_() + index = Tensor.zeros(4, 4, dtype=dtypes.uint8) + self.assertRaisesRegex(IndexError, 'mask', lambda: arr[index]) + self.assertRaisesRegex(IndexError, 'mask', lambda: arr[(slice(None), index)]) + ''' -# def test_boolean_indexing_onedim(self): -# # Indexing a 2-dimensional array with -# # boolean array of length one -# a = tensor([[0., 0., 0.]]) -# b = tensor([True]) -# numpy_testing_assert_equal_helper(a[b], a) -# # boolean assignment -# a[b] = 1. -# numpy_testing_assert_equal_helper(a, tensor([[1., 1., 1.]])) + # TODO setitem + # TODO bool indexing + ''' + def test_boolean_indexing_onedim(self): + # Indexing a 2-dimensional array with + # boolean array of length one + a = Tensor([[0., 0., 0.]]) + b = Tensor([True]) + numpy_testing_assert_equal_helper(a[b], a) + # boolean assignment + a[b] = 1. + numpy_testing_assert_equal_helper(a, Tensor([[1., 1., 1.]])) + ''' -# def test_boolean_assignment_value_mismatch(self): -# # A boolean assignment should fail when the shape of the values -# # cannot be broadcast to the subscription. (see also gh-3458) -# a = torch.arange(0, 4) + # TODO setitem + # TODO bool indexing + ''' + def test_boolean_assignment_value_mismatch(self): + # A boolean assignment should fail when the shape of the values + # cannot be broadcast to the subscription. (see also gh-3458) + a = Tensor.arange(0, 4) -# def f(a, v): -# a[a > -1] = tensor(v).to(device) + def f(a, v): + a[a > -1] = Tensor(v) -# self.assertRaisesRegex(Exception, 'shape mismatch', f, a, []) -# self.assertRaisesRegex(Exception, 'shape mismatch', f, a, [1, 2, 3]) -# self.assertRaisesRegex(Exception, 'shape mismatch', f, a[:1], [1, 2, 3]) + self.assertRaisesRegex(Exception, 'shape mismatch', f, a, []) + self.assertRaisesRegex(Exception, 'shape mismatch', f, a, [1, 2, 3]) + self.assertRaisesRegex(Exception, 'shape mismatch', f, a[:1], [1, 2, 3]) + ''' -# def test_boolean_indexing_twodim(self): -# # Indexing a 2-dimensional array with -# # 2-dimensional boolean array -# a = tensor([[1, 2, 3], -# [4, 5, 6], -# [7, 8, 9]]) -# b = tensor([[True, False, True], -# [False, True, False], -# [True, False, True]]) -# numpy_testing_assert_equal_helper(a[b], tensor([1, 3, 5, 7, 9])) -# numpy_testing_assert_equal_helper(a[b[1]], tensor([[4, 5, 6]])) -# numpy_testing_assert_equal_helper(a[b[0]], a[b[2]]) + # TODO setitem + # TODO bool indexing + ''' + def test_boolean_indexing_twodim(self): + # Indexing a 2-dimensional array with + # 2-dimensional boolean array + a = Tensor([[1, 2, 3], + [4, 5, 6], + [7, 8, 9]]) + b = Tensor([[True, False, True], + [False, True, False], + [True, False, True]]) + numpy_testing_assert_equal_helper(a[b], Tensor([1, 3, 5, 7, 9])) + numpy_testing_assert_equal_helper(a[b[1]], Tensor([[4, 5, 6]])) + numpy_testing_assert_equal_helper(a[b[0]], a[b[2]]) -# # boolean assignment -# a[b] = 0 -# numpy_testing_assert_equal_helper(a, tensor([[0, 2, 0], -# [4, 0, 6], -# [0, 8, 0]])) + # boolean assignment + a[b] = 0 + numpy_testing_assert_equal_helper(a, Tensor([[0, 2, 0], + [4, 0, 6], + [0, 8, 0]])) + ''' -# def test_boolean_indexing_weirdness(self): -# # Weird boolean indexing things -# a = torch.ones((2, 3, 4)) -# numpy_testing_assert_equal_helper((0, 2, 3, 4), a[False, True, ...].shape) -# numpy_testing_assert_equal_helper(torch.ones(1, 2), a[True, [0, 1], True, True, [1], [[2]]]) -# self.assertRaises(IndexError, lambda: a[False, [0, 1], ...]) + # TODO bool indexing + ''' + def test_boolean_indexing_weirdness(self): + # Weird boolean indexing things + a = Tensor.ones((2, 3, 4)) + numpy_testing_assert_equal_helper((0, 2, 3, 4), a[False, True, ...].shape) + numpy_testing_assert_equal_helper(Tensor.ones(1, 2), a[True, [0, 1], True, True, [1], [[2]]]) + self.assertRaises(IndexError, lambda: a[False, [0, 1], ...]) + ''' -# def test_boolean_indexing_weirdness_tensors(self): -# # Weird boolean indexing things -# false = np.array(False) -# true = np.array(True) -# a = torch.ones((2, 3, 4)) -# numpy_testing_assert_equal_helper((0, 2, 3, 4), a[False, True, ...].shape) -# numpy_testing_assert_equal_helper(torch.ones(1, 2), a[true, [0, 1], true, true, [1], [[2]]]) -# self.assertRaises(IndexError, lambda: a[false, [0, 1], ...]) + # TODO bool indexing + ''' + def test_boolean_indexing_weirdness_tensors(self): + # Weird boolean indexing things + false = np.array(False) + true = np.array(True) + a = torch.ones((2, 3, 4)) + numpy_testing_assert_equal_helper((0, 2, 3, 4), a[False, True, ...].shape) + numpy_testing_assert_equal_helper(torch.ones(1, 2), a[true, [0, 1], true, true, [1], [[2]]]) + self.assertRaises(IndexError, lambda: a[false, [0, 1], ...]) + ''' -# def test_boolean_indexing_alldims(self): -# true = np.array(True) -# a = torch.ones((2, 3)) -# numpy_testing_assert_equal_helper((1, 2, 3), a[True, True].shape) -# numpy_testing_assert_equal_helper((1, 2, 3), a[true, true].shape) + # TODO bool indexing + ''' + def test_boolean_indexing_alldims(self): + true = Tensor(True) + a = Tensor.ones((2, 3)) + numpy_testing_assert_equal_helper((1, 2, 3), a[True, True].shape) + numpy_testing_assert_equal_helper((1, 2, 3), a[true, true].shape) + ''' -# def test_boolean_list_indexing(self): -# # Indexing a 2-dimensional array with -# # boolean lists -# a = tensor([[1, 2, 3], -# [4, 5, 6], -# [7, 8, 9]]) -# b = [True, False, False] -# c = [True, True, False] -# numpy_testing_assert_equal_helper(a[b], tensor([[1, 2, 3]])) -# numpy_testing_assert_equal_helper(a[b, b], tensor([1])) -# numpy_testing_assert_equal_helper(a[c], tensor([[1, 2, 3], [4, 5, 6]])) -# numpy_testing_assert_equal_helper(a[c, c], tensor([1, 5])) + # TODO bool indexing + # NOTE this is easier + ''' + def test_boolean_list_indexing(self): + # Indexing a 2-dimensional array with + # boolean lists + a = Tensor([[1, 2, 3], + [4, 5, 6], + [7, 8, 9]]) + b = [True, False, False] + c = [True, True, False] + numpy_testing_assert_equal_helper(a[b], Tensor([[1, 2, 3]])) + numpy_testing_assert_equal_helper(a[b, b], Tensor([1])) + numpy_testing_assert_equal_helper(a[c], Tensor([[1, 2, 3], [4, 5, 6]])) + numpy_testing_assert_equal_helper(a[c, c], Tensor([1, 5])) + ''' -# def test_everything_returns_views(self): -# # Before `...` would return a itself. -# a = tensor([5]) + def test_everything_returns_views(self): + # Before `...` would return a itself. + a = Tensor([5]) -# self.assertIsNot(a, a[()]) -# self.assertIsNot(a, a[...]) -# self.assertIsNot(a, a[:]) + self.assertIsNot(a, a[()]) + self.assertIsNot(a, a[...]) + self.assertIsNot(a, a[:]) -# def test_broaderrors_indexing(self): -# a = torch.zeros(5, 5) -# self.assertRaisesRegex(IndexError, 'shape mismatch', a.__getitem__, ([0, 1], [0, 1, 2])) -# self.assertRaisesRegex(IndexError, 'shape mismatch', a.__setitem__, ([0, 1], [0, 1, 2]), 0) + # TODO shape mismatch fancy indexing error + ''' + def test_broaderrors_indexing(self): + a = Tensor.zeros(5, 5) + self.assertRaisesRegex(IndexError, 'shape mismatch', a.__getitem__, ([0, 1], [0, 1, 2])) + self.assertRaisesRegex(IndexError, 'shape mismatch', a.__setitem__, ([0, 1], [0, 1, 2]), 0) + ''' -# def test_trivial_fancy_out_of_bounds(self): -# a = torch.zeros(5) -# ind = torch.ones(20, dtype=torch.int64) -# if a.is_cuda: -# raise unittest.SkipTest('CUDA asserts instead of raising an exception') -# ind[-1] = 10 -# self.assertRaises(IndexError, a.__getitem__, ind) -# self.assertRaises(IndexError, a.__setitem__, ind, 0) -# ind = torch.ones(20, dtype=torch.int64) -# ind[0] = 11 -# self.assertRaises(IndexError, a.__getitem__, ind) -# self.assertRaises(IndexError, a.__setitem__, ind, 0) + # TODO setitem + ''' + def test_trivial_fancy_out_of_bounds(self): + a = Tensor.zeros(5) + ind = Tensor.ones(20, dtype=dtypes.int64) + ind[-1] = 10 + self.assertRaises(IndexError, a.__getitem__, ind) + self.assertRaises(IndexError, a.__setitem__, ind, 0) + ind = Tensor.ones(20, dtype=dtypes.int64) + ind[0] = 11 + self.assertRaises(IndexError, a.__getitem__, ind) + self.assertRaises(IndexError, a.__setitem__, ind, 0) + ''' -# def test_index_is_larger(self): -# # Simple case of fancy index broadcasting of the index. -# a = torch.zeros((5, 5)) -# a[[[0], [1], [2]], [0, 1, 2]] = tensor([2., 3., 4.]) + # TODO setitem + ''' + def test_index_is_larger(self): + # Simple case of fancy index broadcasting of the index. + a = Tensor.zeros((5, 5)) + a[[[0], [1], [2]], [0, 1, 2]] = Tensor([2., 3., 4.]) -# self.assertTrue((a[:3, :3] == tensor([2., 3., 4.])).all()) + self.assertTrue((a[:3, :3] == all_(Tensor([2., 3., 4.])))) + ''' -# def test_broadcast_subspace(self): -# a = torch.zeros((100, 100)) -# v = torch.arange(0., 100)[:, None] -# b = torch.arange(99, -1, -1).long() -# a[b] = v -# expected = b.float().unsqueeze(1).expand(100, 100) -# numpy_testing_assert_equal_helper(a, expected) - -# def test_truncate_leading_1s(self): -# col_max = torch.randn(1, 4) -# kernel = col_max.T * col_max # [4, 4] tensor -# kernel2 = kernel.clone() -# # Set the diagonal -# kernel[range(len(kernel)), range(len(kernel))] = torch.square(col_max) -# torch.diagonal(kernel2).copy_(torch.square(col_max.view(4))) -# numpy_testing_assert_equal_helper(kernel, kernel2) + # TODO setitem + ''' + def test_broadcast_subspace(self): + a = Tensor.zeros((100, 100)) + v = Tensor.arange(0., 100)[:, None] + b = Tensor.arange(99, -1, -1).cast(dtypes.int64) + a[b] = v + expected = b.float().unsqueeze(1).expand(100, 100) + numpy_testing_assert_equal_helper(a, expected) + ''' + # TODO copy_ + ''' + def test_truncate_leading_1s(self): + col_max = Tensor.randn(1, 4) + kernel = col_max.T * col_max # [4, 4] tensor + kernel2 = clone(kernel) + # Set the diagonal + kernel[range(len(kernel)), range(len(kernel))] = col_max.square() + kernel2 = diagonal(kernel2) + # torch.diagonal(kernel2).copy_(torch.square(col_max.view(4))) + kernel2 = copy_(kernel2, col_max.reshape(4).square()) + numpy_testing_assert_equal_helper(kernel, kernel2) + ''' if __name__ == '__main__': unittest.main() \ No newline at end of file diff --git a/test/test_ops.py b/test/test_ops.py index e0481b837a..8f4c36097f 100644 --- a/test/test_ops.py +++ b/test/test_ops.py @@ -666,11 +666,13 @@ class TestOps(unittest.TestCase): def test_slice_errors(self): a = Tensor.ones(4, 3) + b = Tensor(2) with self.assertRaises(IndexError): a[1, 77, 77, 77] # IndexError: (finds too many indices before the out of bounds) with self.assertRaises(IndexError): a[1, 77] # IndexError: (out of bounds). with self.assertRaises(IndexError): a[1, -77] with self.assertRaises(IndexError): a[..., ...] # IndexError: only single ellipsis with self.assertRaises(ValueError): a[::0, 1] # no 0 strides + with self.assertRaises(IndexError): b[:] # slice cannot be applied to a 0-dim tensor def test_slice_ellipsis(self): helper_test_op([(3,3,3,3)], lambda x: x[..., 0], lambda x: x[..., 0]) @@ -1312,8 +1314,7 @@ class TestOps(unittest.TestCase): # TODO: currently we do not support tensor indexing for list of list tensor # ex: torch.tensor([1,2])[[[[torch.tensor(1)]]]] -> tensor([[2]]) # currently we return ValueError: setting an array element with a sequence. - # E TypeError: only integer tensors of a single element can be converted to an index - + # TypeError: only integer tensors of a single element can be converted to an index def test_gather(self): # indices cannot have gradient diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index e2c5f9865d..8f0cad1a77 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -313,8 +313,8 @@ class Tensor: # 1. indices normalization and validation # treat internal tuples and lists as Tensors and standardize indices to list type if isinstance(indices, (tuple, list)): - if isinstance(indices, list) and all(isinstance(i, int) for i in indices): indices = [Tensor(indices)] # special case , a lil ugly - else: indices = [Tensor(list(i)) if isinstance(i, (tuple, list)) else i for i in indices] + if isinstance(indices, list) and all(isinstance(i, int) for i in indices): indices = [Tensor(indices, dtype=dtypes.int32, requires_grad=False, device=self.device)] # special case , a lil ugly + else: indices = [Tensor(list(i), dtype=dtypes.int32, requires_grad=False, device=self.device) if isinstance(i, (tuple, list)) else i for i in indices] else: indices = [indices] # filter ellipsis and fill with slice(None) or fill rest of indices with slice(None) @@ -334,6 +334,7 @@ class Tensor: for dim,i in enumerate(indices_filtered): type_dim[type(i)].append(dim) # validation! raise Errors + if slice in type_dim and self.ndim == 0: raise IndexError("slice cannot be applied to a 0-dim tensor.") if len(ellipsis_idx) > 1: raise IndexError("an index can only have a single ellipsis ('...')") if float in type_dim: raise IndexError("float type is not valid index") if any(isinstance(i, slice) and i.step == 0 for i in indices): raise ValueError('slice step cannot be 0')