From c281d3a0cb435792c88f5015fafd33af40e2004e Mon Sep 17 00:00:00 2001 From: Jianping Wu Date: Wed, 1 May 2024 23:44:53 +0000 Subject: [PATCH] Enable UFMT on test_indexing&test_view_ops (#125112) Part of https://github.com/pytorch/pytorch/issues/123062 Pull Request resolved: https://github.com/pytorch/pytorch/pull/125112 Approved by: https://github.com/ezyang --- .lintrunner.toml | 2 - test/test_indexing.py | 877 ++++++++++++++++++++++++++---------------- test/test_view_ops.py | 555 +++++++++++++++++--------- 3 files changed, 933 insertions(+), 501 deletions(-) diff --git a/.lintrunner.toml b/.lintrunner.toml index a945814b23af..21dfdf5b3f10 100644 --- a/.lintrunner.toml +++ b/.lintrunner.toml @@ -1071,7 +1071,6 @@ exclude_patterns = [ 'test/test_fx_reinplace_pass.py', 'test/test_hub.py', 'test/test_import_stats.py', - 'test/test_indexing.py', 'test/test_itt.py', 'test/test_jit.py', 'test/test_jit_autocast.py', @@ -1157,7 +1156,6 @@ exclude_patterns = [ 'test/test_type_promotion.py', 'test/test_unary_ufuncs.py', 'test/test_utils.py', - 'test/test_view_ops.py', 'test/test_vulkan.py', 'test/test_xnnpack_integration.py', 'test/torch_np/numpy_test/**/*.py', diff --git a/test/test_indexing.py b/test/test_indexing.py index 195a3144c33c..f34fa4c5669b 100644 --- a/test/test_indexing.py +++ b/test/test_indexing.py @@ -1,27 +1,38 @@ # Owner(s): ["module: tests"] -import torch -from torch import tensor +import operator +import random import unittest import warnings -import random from functools import reduce import numpy as np +import torch +from torch import tensor from torch.testing import make_tensor -from torch.testing._internal.common_utils import ( - TestCase, run_tests, skipIfTorchDynamo, DeterministicGuard, serialTest, TEST_CUDA) from torch.testing._internal.common_device_type import ( - instantiate_device_type_tests, onlyCUDA, dtypes, dtypesIfCPU, dtypesIfCUDA, - onlyNativeDeviceTypes, skipXLA) -import operator + dtypes, + dtypesIfCPU, + dtypesIfCUDA, + instantiate_device_type_tests, + onlyCUDA, + onlyNativeDeviceTypes, + skipXLA, +) +from torch.testing._internal.common_utils import ( + DeterministicGuard, + run_tests, + serialTest, + skipIfTorchDynamo, + TEST_CUDA, + TestCase, +) class TestIndexing(TestCase): def test_index(self, device): - def consec(size, start=1): sequence = torch.ones(torch.tensor(size).prod(0)).cumsum(0) sequence.add_(start - 1) @@ -30,7 +41,9 @@ class TestIndexing(TestCase): reference = consec((3, 3, 3)).to(device) # empty tensor indexing - self.assertEqual(reference[torch.LongTensor().to(device)], reference.new(0, 3, 3)) + self.assertEqual( + reference[torch.LongTensor().to(device)], reference.new(0, 3, 3) + ) self.assertEqual(reference[0], consec((3, 3)), atol=0, rtol=0) self.assertEqual(reference[1], consec((3, 3), 10), atol=0, rtol=0) @@ -41,10 +54,15 @@ class TestIndexing(TestCase): self.assertEqual(reference[:], consec((3, 3, 3)), atol=0, rtol=0) # indexing with Ellipsis - self.assertEqual(reference[..., 2], torch.tensor([[3., 6., 9.], - [12., 15., 18.], - [21., 24., 27.]]), atol=0, rtol=0) - self.assertEqual(reference[0, ..., 2], torch.tensor([3., 6., 9.]), atol=0, rtol=0) + self.assertEqual( + reference[..., 2], + torch.tensor([[3.0, 6.0, 9.0], [12.0, 15.0, 18.0], [21.0, 24.0, 27.0]]), + atol=0, + rtol=0, + ) + self.assertEqual( + reference[0, ..., 2], torch.tensor([3.0, 6.0, 9.0]), atol=0, rtol=0 + ) self.assertEqual(reference[..., 2], reference[:, :, 2], atol=0, rtol=0) self.assertEqual(reference[0, ..., 2], reference[0, :, 2], atol=0, rtol=0) self.assertEqual(reference[0, 2, ...], reference[0, 2], atol=0, rtol=0) @@ -55,9 +73,15 @@ class TestIndexing(TestCase): self.assertEqual(reference[...], reference, atol=0, rtol=0) reference_5d = consec((3, 3, 3, 3, 3)).to(device) - self.assertEqual(reference_5d[..., 1, 0], reference_5d[:, :, :, 1, 0], atol=0, rtol=0) - self.assertEqual(reference_5d[2, ..., 1, 0], reference_5d[2, :, :, 1, 0], atol=0, rtol=0) - self.assertEqual(reference_5d[2, 1, 0, ..., 1], reference_5d[2, 1, 0, :, 1], atol=0, rtol=0) + self.assertEqual( + reference_5d[..., 1, 0], reference_5d[:, :, :, 1, 0], atol=0, rtol=0 + ) + self.assertEqual( + reference_5d[2, ..., 1, 0], reference_5d[2, :, :, 1, 0], atol=0, rtol=0 + ) + self.assertEqual( + reference_5d[2, 1, 0, ..., 1], reference_5d[2, 1, 0, :, 1], atol=0, rtol=0 + ) self.assertEqual(reference_5d[...], reference_5d, atol=0, rtol=0) # LongTensor indexing @@ -70,10 +94,18 @@ class TestIndexing(TestCase): # None indexing self.assertEqual(reference[2, None], reference[2].unsqueeze(0)) - self.assertEqual(reference[2, None, None], reference[2].unsqueeze(0).unsqueeze(0)) + self.assertEqual( + reference[2, None, None], reference[2].unsqueeze(0).unsqueeze(0) + ) self.assertEqual(reference[2:4, None], reference[2:4].unsqueeze(1)) - self.assertEqual(reference[None, 2, None, None], reference.unsqueeze(0)[:, 2].unsqueeze(0).unsqueeze(0)) - self.assertEqual(reference[None, 2:5, None, None], reference.unsqueeze(0)[:, 2:5].unsqueeze(2).unsqueeze(2)) + self.assertEqual( + reference[None, 2, None, None], + reference.unsqueeze(0)[:, 2].unsqueeze(0).unsqueeze(0), + ) + self.assertEqual( + reference[None, 2:5, None, None], + reference.unsqueeze(0)[:, 2:5].unsqueeze(2).unsqueeze(2), + ) # indexing 0-length slice self.assertEqual(torch.empty(0, 5, 5), reference[slice(0)]) @@ -84,13 +116,28 @@ class TestIndexing(TestCase): # indexing with step reference = consec((10, 10, 10)).to(device) self.assertEqual(reference[1:5:2], torch.stack([reference[1], reference[3]], 0)) - self.assertEqual(reference[1:6:2], torch.stack([reference[1], reference[3], reference[5]], 0)) + self.assertEqual( + reference[1:6:2], torch.stack([reference[1], reference[3], reference[5]], 0) + ) self.assertEqual(reference[1:9:4], torch.stack([reference[1], reference[5]], 0)) - self.assertEqual(reference[2:4, 1:5:2], torch.stack([reference[2:4, 1], reference[2:4, 3]], 1)) - self.assertEqual(reference[3, 1:6:2], torch.stack([reference[3, 1], reference[3, 3], reference[3, 5]], 0)) - self.assertEqual(reference[None, 2, 1:9:4], torch.stack([reference[2, 1], reference[2, 5]], 0).unsqueeze(0)) - self.assertEqual(reference[:, 2, 1:6:2], - torch.stack([reference[:, 2, 1], reference[:, 2, 3], reference[:, 2, 5]], 1)) + self.assertEqual( + reference[2:4, 1:5:2], + torch.stack([reference[2:4, 1], reference[2:4, 3]], 1), + ) + self.assertEqual( + reference[3, 1:6:2], + torch.stack([reference[3, 1], reference[3, 3], reference[3, 5]], 0), + ) + self.assertEqual( + reference[None, 2, 1:9:4], + torch.stack([reference[2, 1], reference[2, 5]], 0).unsqueeze(0), + ) + self.assertEqual( + reference[:, 2, 1:6:2], + torch.stack( + [reference[:, 2, 1], reference[:, 2, 3], reference[:, 2, 5]], 1 + ), + ) lst = [list(range(i, i + 10)) for i in range(0, 100, 10)] tensor = torch.DoubleTensor(lst).to(device) @@ -156,23 +203,33 @@ class TestIndexing(TestCase): def validate_indexing(x): self.assertEqual(x[[0]], consec((1,))) - self.assertEqual(x[ri([0]), ], consec((1,))) - self.assertEqual(x[ri([3]), ], consec((1,), 4)) + self.assertEqual(x[ri([0]),], consec((1,))) + self.assertEqual(x[ri([3]),], consec((1,), 4)) self.assertEqual(x[[2, 3, 4]], consec((3,), 3)) - self.assertEqual(x[ri([2, 3, 4]), ], consec((3,), 3)) - self.assertEqual(x[ri([0, 2, 4]), ], torch.tensor([1, 3, 5], dtype=dtype, device=device)) + self.assertEqual(x[ri([2, 3, 4]),], consec((3,), 3)) + self.assertEqual( + x[ri([0, 2, 4]),], torch.tensor([1, 3, 5], dtype=dtype, device=device) + ) def validate_setting(x): x[[0]] = -2 self.assertEqual(x[[0]], torch.tensor([-2], dtype=dtype, device=device)) x[[0]] = -1 - self.assertEqual(x[ri([0]), ], torch.tensor([-1], dtype=dtype, device=device)) + self.assertEqual( + x[ri([0]),], torch.tensor([-1], dtype=dtype, device=device) + ) x[[2, 3, 4]] = 4 - self.assertEqual(x[[2, 3, 4]], torch.tensor([4, 4, 4], dtype=dtype, device=device)) - x[ri([2, 3, 4]), ] = 3 - self.assertEqual(x[ri([2, 3, 4]), ], torch.tensor([3, 3, 3], dtype=dtype, device=device)) - x[ri([0, 2, 4]), ] = torch.tensor([5, 4, 3], dtype=dtype, device=device) - self.assertEqual(x[ri([0, 2, 4]), ], torch.tensor([5, 4, 3], dtype=dtype, device=device)) + self.assertEqual( + x[[2, 3, 4]], torch.tensor([4, 4, 4], dtype=dtype, device=device) + ) + x[ri([2, 3, 4]),] = 3 + self.assertEqual( + x[ri([2, 3, 4]),], torch.tensor([3, 3, 3], dtype=dtype, device=device) + ) + x[ri([0, 2, 4]),] = torch.tensor([5, 4, 3], dtype=dtype, device=device) + self.assertEqual( + x[ri([0, 2, 4]),], torch.tensor([5, 4, 3], dtype=dtype, device=device) + ) # Only validates indexing and setting for halfs if dtype == torch.half: @@ -192,208 +249,300 @@ class TestIndexing(TestCase): # strided is [1, 3, 5, 7] reference = consec((10,)) strided = torch.tensor((), dtype=dtype, device=device) - strided.set_(reference.storage(), storage_offset=0, - size=torch.Size([4]), stride=[2]) + strided.set_( + reference.storage(), storage_offset=0, size=torch.Size([4]), stride=[2] + ) self.assertEqual(strided[[0]], torch.tensor([1], dtype=dtype, device=device)) - self.assertEqual(strided[ri([0]), ], torch.tensor([1], dtype=dtype, device=device)) - self.assertEqual(strided[ri([3]), ], torch.tensor([7], dtype=dtype, device=device)) - self.assertEqual(strided[[1, 2]], torch.tensor([3, 5], dtype=dtype, device=device)) - self.assertEqual(strided[ri([1, 2]), ], torch.tensor([3, 5], dtype=dtype, device=device)) - self.assertEqual(strided[ri([[2, 1], [0, 3]]), ], - torch.tensor([[5, 3], [1, 7]], dtype=dtype, device=device)) + self.assertEqual( + strided[ri([0]),], torch.tensor([1], dtype=dtype, device=device) + ) + self.assertEqual( + strided[ri([3]),], torch.tensor([7], dtype=dtype, device=device) + ) + self.assertEqual( + strided[[1, 2]], torch.tensor([3, 5], dtype=dtype, device=device) + ) + self.assertEqual( + strided[ri([1, 2]),], torch.tensor([3, 5], dtype=dtype, device=device) + ) + self.assertEqual( + strided[ri([[2, 1], [0, 3]]),], + torch.tensor([[5, 3], [1, 7]], dtype=dtype, device=device), + ) # stride is [4, 8] strided = torch.tensor((), dtype=dtype, device=device) - strided.set_(reference.storage(), storage_offset=4, - size=torch.Size([2]), stride=[4]) + strided.set_( + reference.storage(), storage_offset=4, size=torch.Size([2]), stride=[4] + ) self.assertEqual(strided[[0]], torch.tensor([5], dtype=dtype, device=device)) - self.assertEqual(strided[ri([0]), ], torch.tensor([5], dtype=dtype, device=device)) - self.assertEqual(strided[ri([1]), ], torch.tensor([9], dtype=dtype, device=device)) - self.assertEqual(strided[[0, 1]], torch.tensor([5, 9], dtype=dtype, device=device)) - self.assertEqual(strided[ri([0, 1]), ], torch.tensor([5, 9], dtype=dtype, device=device)) - self.assertEqual(strided[ri([[0, 1], [1, 0]]), ], - torch.tensor([[5, 9], [9, 5]], dtype=dtype, device=device)) + self.assertEqual( + strided[ri([0]),], torch.tensor([5], dtype=dtype, device=device) + ) + self.assertEqual( + strided[ri([1]),], torch.tensor([9], dtype=dtype, device=device) + ) + self.assertEqual( + strided[[0, 1]], torch.tensor([5, 9], dtype=dtype, device=device) + ) + self.assertEqual( + strided[ri([0, 1]),], torch.tensor([5, 9], dtype=dtype, device=device) + ) + self.assertEqual( + strided[ri([[0, 1], [1, 0]]),], + torch.tensor([[5, 9], [9, 5]], dtype=dtype, device=device), + ) # reference is 1 2 # 3 4 # 5 6 reference = consec((3, 2)) - self.assertEqual(reference[ri([0, 1, 2]), ri([0])], torch.tensor([1, 3, 5], dtype=dtype, device=device)) - self.assertEqual(reference[ri([0, 1, 2]), ri([1])], torch.tensor([2, 4, 6], dtype=dtype, device=device)) + self.assertEqual( + reference[ri([0, 1, 2]), ri([0])], + torch.tensor([1, 3, 5], dtype=dtype, device=device), + ) + self.assertEqual( + reference[ri([0, 1, 2]), ri([1])], + torch.tensor([2, 4, 6], dtype=dtype, device=device), + ) self.assertEqual(reference[ri([0]), ri([0])], consec((1,))) self.assertEqual(reference[ri([2]), ri([1])], consec((1,), 6)) - self.assertEqual(reference[[ri([0, 0]), ri([0, 1])]], torch.tensor([1, 2], dtype=dtype, device=device)) - self.assertEqual(reference[[ri([0, 1, 1, 0, 2]), ri([1])]], - torch.tensor([2, 4, 4, 2, 6], dtype=dtype, device=device)) - self.assertEqual(reference[[ri([0, 0, 1, 1]), ri([0, 1, 0, 0])]], - torch.tensor([1, 2, 3, 3], dtype=dtype, device=device)) + self.assertEqual( + reference[[ri([0, 0]), ri([0, 1])]], + torch.tensor([1, 2], dtype=dtype, device=device), + ) + self.assertEqual( + reference[[ri([0, 1, 1, 0, 2]), ri([1])]], + torch.tensor([2, 4, 4, 2, 6], dtype=dtype, device=device), + ) + self.assertEqual( + reference[[ri([0, 0, 1, 1]), ri([0, 1, 0, 0])]], + torch.tensor([1, 2, 3, 3], dtype=dtype, device=device), + ) - rows = ri([[0, 0], - [1, 2]]) - columns = [0], - self.assertEqual(reference[rows, columns], torch.tensor([[1, 1], - [3, 5]], dtype=dtype, device=device)) + rows = ri([[0, 0], [1, 2]]) + columns = ([0],) + self.assertEqual( + reference[rows, columns], + torch.tensor([[1, 1], [3, 5]], dtype=dtype, device=device), + ) - rows = ri([[0, 0], - [1, 2]]) + rows = ri([[0, 0], [1, 2]]) columns = ri([1, 0]) - self.assertEqual(reference[rows, columns], torch.tensor([[2, 1], - [4, 5]], dtype=dtype, device=device)) - rows = ri([[0, 0], - [1, 2]]) - columns = ri([[0, 1], - [1, 0]]) - self.assertEqual(reference[rows, columns], torch.tensor([[1, 2], - [4, 5]], dtype=dtype, device=device)) + self.assertEqual( + reference[rows, columns], + torch.tensor([[2, 1], [4, 5]], dtype=dtype, device=device), + ) + rows = ri([[0, 0], [1, 2]]) + columns = ri([[0, 1], [1, 0]]) + self.assertEqual( + reference[rows, columns], + torch.tensor([[1, 2], [4, 5]], dtype=dtype, device=device), + ) # setting values reference[ri([0]), ri([1])] = -1 - self.assertEqual(reference[ri([0]), ri([1])], torch.tensor([-1], dtype=dtype, device=device)) - reference[ri([0, 1, 2]), ri([0])] = torch.tensor([-1, 2, -4], dtype=dtype, device=device) - self.assertEqual(reference[ri([0, 1, 2]), ri([0])], - torch.tensor([-1, 2, -4], dtype=dtype, device=device)) - reference[rows, columns] = torch.tensor([[4, 6], [2, 3]], dtype=dtype, device=device) - self.assertEqual(reference[rows, columns], - torch.tensor([[4, 6], [2, 3]], dtype=dtype, device=device)) + self.assertEqual( + reference[ri([0]), ri([1])], torch.tensor([-1], dtype=dtype, device=device) + ) + reference[ri([0, 1, 2]), ri([0])] = torch.tensor( + [-1, 2, -4], dtype=dtype, device=device + ) + self.assertEqual( + reference[ri([0, 1, 2]), ri([0])], + torch.tensor([-1, 2, -4], dtype=dtype, device=device), + ) + reference[rows, columns] = torch.tensor( + [[4, 6], [2, 3]], dtype=dtype, device=device + ) + self.assertEqual( + reference[rows, columns], + torch.tensor([[4, 6], [2, 3]], dtype=dtype, device=device), + ) # Verify still works with Transposed (i.e. non-contiguous) Tensors - reference = torch.tensor([[0, 1, 2, 3], - [4, 5, 6, 7], - [8, 9, 10, 11]], dtype=dtype, device=device).t_() + reference = torch.tensor( + [[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11]], dtype=dtype, device=device + ).t_() # Transposed: [[0, 4, 8], # [1, 5, 9], # [2, 6, 10], # [3, 7, 11]] - self.assertEqual(reference[ri([0, 1, 2]), ri([0])], - torch.tensor([0, 1, 2], dtype=dtype, device=device)) - self.assertEqual(reference[ri([0, 1, 2]), ri([1])], - torch.tensor([4, 5, 6], dtype=dtype, device=device)) - self.assertEqual(reference[ri([0]), ri([0])], - torch.tensor([0], dtype=dtype, device=device)) - self.assertEqual(reference[ri([2]), ri([1])], - torch.tensor([6], dtype=dtype, device=device)) - self.assertEqual(reference[[ri([0, 0]), ri([0, 1])]], - torch.tensor([0, 4], dtype=dtype, device=device)) - self.assertEqual(reference[[ri([0, 1, 1, 0, 3]), ri([1])]], - torch.tensor([4, 5, 5, 4, 7], dtype=dtype, device=device)) - self.assertEqual(reference[[ri([0, 0, 1, 1]), ri([0, 1, 0, 0])]], - torch.tensor([0, 4, 1, 1], dtype=dtype, device=device)) + self.assertEqual( + reference[ri([0, 1, 2]), ri([0])], + torch.tensor([0, 1, 2], dtype=dtype, device=device), + ) + self.assertEqual( + reference[ri([0, 1, 2]), ri([1])], + torch.tensor([4, 5, 6], dtype=dtype, device=device), + ) + self.assertEqual( + reference[ri([0]), ri([0])], torch.tensor([0], dtype=dtype, device=device) + ) + self.assertEqual( + reference[ri([2]), ri([1])], torch.tensor([6], dtype=dtype, device=device) + ) + self.assertEqual( + reference[[ri([0, 0]), ri([0, 1])]], + torch.tensor([0, 4], dtype=dtype, device=device), + ) + self.assertEqual( + reference[[ri([0, 1, 1, 0, 3]), ri([1])]], + torch.tensor([4, 5, 5, 4, 7], dtype=dtype, device=device), + ) + self.assertEqual( + reference[[ri([0, 0, 1, 1]), ri([0, 1, 0, 0])]], + torch.tensor([0, 4, 1, 1], dtype=dtype, device=device), + ) - rows = ri([[0, 0], - [1, 2]]) - columns = [0], - self.assertEqual(reference[rows, columns], - torch.tensor([[0, 0], [1, 2]], dtype=dtype, device=device)) + rows = ri([[0, 0], [1, 2]]) + columns = ([0],) + self.assertEqual( + reference[rows, columns], + torch.tensor([[0, 0], [1, 2]], dtype=dtype, device=device), + ) - rows = ri([[0, 0], - [1, 2]]) + rows = ri([[0, 0], [1, 2]]) columns = ri([1, 0]) - self.assertEqual(reference[rows, columns], - torch.tensor([[4, 0], [5, 2]], dtype=dtype, device=device)) - rows = ri([[0, 0], - [1, 3]]) - columns = ri([[0, 1], - [1, 2]]) - self.assertEqual(reference[rows, columns], - torch.tensor([[0, 4], [5, 11]], dtype=dtype, device=device)) + self.assertEqual( + reference[rows, columns], + torch.tensor([[4, 0], [5, 2]], dtype=dtype, device=device), + ) + rows = ri([[0, 0], [1, 3]]) + columns = ri([[0, 1], [1, 2]]) + self.assertEqual( + reference[rows, columns], + torch.tensor([[0, 4], [5, 11]], dtype=dtype, device=device), + ) # setting values reference[ri([0]), ri([1])] = -1 - self.assertEqual(reference[ri([0]), ri([1])], - torch.tensor([-1], dtype=dtype, device=device)) - reference[ri([0, 1, 2]), ri([0])] = torch.tensor([-1, 2, -4], dtype=dtype, device=device) - self.assertEqual(reference[ri([0, 1, 2]), ri([0])], - torch.tensor([-1, 2, -4], dtype=dtype, device=device)) - reference[rows, columns] = torch.tensor([[4, 6], [2, 3]], dtype=dtype, device=device) - self.assertEqual(reference[rows, columns], - torch.tensor([[4, 6], [2, 3]], dtype=dtype, device=device)) + self.assertEqual( + reference[ri([0]), ri([1])], torch.tensor([-1], dtype=dtype, device=device) + ) + reference[ri([0, 1, 2]), ri([0])] = torch.tensor( + [-1, 2, -4], dtype=dtype, device=device + ) + self.assertEqual( + reference[ri([0, 1, 2]), ri([0])], + torch.tensor([-1, 2, -4], dtype=dtype, device=device), + ) + reference[rows, columns] = torch.tensor( + [[4, 6], [2, 3]], dtype=dtype, device=device + ) + self.assertEqual( + reference[rows, columns], + torch.tensor([[4, 6], [2, 3]], dtype=dtype, device=device), + ) # stride != 1 # strided is [[1 3 5 7], # [9 11 13 15]] - reference = torch.arange(0., 24, dtype=dtype, device=device).view(3, 8) + reference = torch.arange(0.0, 24, dtype=dtype, device=device).view(3, 8) strided = torch.tensor((), dtype=dtype, device=device) - strided.set_(reference.storage(), 1, size=torch.Size([2, 4]), - stride=[8, 2]) + strided.set_(reference.storage(), 1, size=torch.Size([2, 4]), stride=[8, 2]) - self.assertEqual(strided[ri([0, 1]), ri([0])], - torch.tensor([1, 9], dtype=dtype, device=device)) - self.assertEqual(strided[ri([0, 1]), ri([1])], - torch.tensor([3, 11], dtype=dtype, device=device)) - self.assertEqual(strided[ri([0]), ri([0])], - torch.tensor([1], dtype=dtype, device=device)) - self.assertEqual(strided[ri([1]), ri([3])], - torch.tensor([15], dtype=dtype, device=device)) - self.assertEqual(strided[[ri([0, 0]), ri([0, 3])]], - torch.tensor([1, 7], dtype=dtype, device=device)) - self.assertEqual(strided[[ri([1]), ri([0, 1, 1, 0, 3])]], - torch.tensor([9, 11, 11, 9, 15], dtype=dtype, device=device)) - self.assertEqual(strided[[ri([0, 0, 1, 1]), ri([0, 1, 0, 0])]], - torch.tensor([1, 3, 9, 9], dtype=dtype, device=device)) + self.assertEqual( + strided[ri([0, 1]), ri([0])], + torch.tensor([1, 9], dtype=dtype, device=device), + ) + self.assertEqual( + strided[ri([0, 1]), ri([1])], + torch.tensor([3, 11], dtype=dtype, device=device), + ) + self.assertEqual( + strided[ri([0]), ri([0])], torch.tensor([1], dtype=dtype, device=device) + ) + self.assertEqual( + strided[ri([1]), ri([3])], torch.tensor([15], dtype=dtype, device=device) + ) + self.assertEqual( + strided[[ri([0, 0]), ri([0, 3])]], + torch.tensor([1, 7], dtype=dtype, device=device), + ) + self.assertEqual( + strided[[ri([1]), ri([0, 1, 1, 0, 3])]], + torch.tensor([9, 11, 11, 9, 15], dtype=dtype, device=device), + ) + self.assertEqual( + strided[[ri([0, 0, 1, 1]), ri([0, 1, 0, 0])]], + torch.tensor([1, 3, 9, 9], dtype=dtype, device=device), + ) - rows = ri([[0, 0], - [1, 1]]) - columns = [0], - self.assertEqual(strided[rows, columns], - torch.tensor([[1, 1], [9, 9]], dtype=dtype, device=device)) + rows = ri([[0, 0], [1, 1]]) + columns = ([0],) + self.assertEqual( + strided[rows, columns], + torch.tensor([[1, 1], [9, 9]], dtype=dtype, device=device), + ) - rows = ri([[0, 1], - [1, 0]]) + rows = ri([[0, 1], [1, 0]]) columns = ri([1, 2]) - self.assertEqual(strided[rows, columns], - torch.tensor([[3, 13], [11, 5]], dtype=dtype, device=device)) - rows = ri([[0, 0], - [1, 1]]) - columns = ri([[0, 1], - [1, 2]]) - self.assertEqual(strided[rows, columns], - torch.tensor([[1, 3], [11, 13]], dtype=dtype, device=device)) + self.assertEqual( + strided[rows, columns], + torch.tensor([[3, 13], [11, 5]], dtype=dtype, device=device), + ) + rows = ri([[0, 0], [1, 1]]) + columns = ri([[0, 1], [1, 2]]) + self.assertEqual( + strided[rows, columns], + torch.tensor([[1, 3], [11, 13]], dtype=dtype, device=device), + ) # setting values # strided is [[10, 11], # [17, 18]] - reference = torch.arange(0., 24, dtype=dtype, device=device).view(3, 8) + reference = torch.arange(0.0, 24, dtype=dtype, device=device).view(3, 8) strided = torch.tensor((), dtype=dtype, device=device) - strided.set_(reference.storage(), 10, size=torch.Size([2, 2]), - stride=[7, 1]) - self.assertEqual(strided[ri([0]), ri([1])], - torch.tensor([11], dtype=dtype, device=device)) + strided.set_(reference.storage(), 10, size=torch.Size([2, 2]), stride=[7, 1]) + self.assertEqual( + strided[ri([0]), ri([1])], torch.tensor([11], dtype=dtype, device=device) + ) strided[ri([0]), ri([1])] = -1 - self.assertEqual(strided[ri([0]), ri([1])], - torch.tensor([-1], dtype=dtype, device=device)) + self.assertEqual( + strided[ri([0]), ri([1])], torch.tensor([-1], dtype=dtype, device=device) + ) - reference = torch.arange(0., 24, dtype=dtype, device=device).view(3, 8) + reference = torch.arange(0.0, 24, dtype=dtype, device=device).view(3, 8) strided = torch.tensor((), dtype=dtype, device=device) - strided.set_(reference.storage(), 10, size=torch.Size([2, 2]), - stride=[7, 1]) - self.assertEqual(strided[ri([0, 1]), ri([1, 0])], - torch.tensor([11, 17], dtype=dtype, device=device)) - strided[ri([0, 1]), ri([1, 0])] = torch.tensor([-1, 2], dtype=dtype, device=device) - self.assertEqual(strided[ri([0, 1]), ri([1, 0])], - torch.tensor([-1, 2], dtype=dtype, device=device)) + strided.set_(reference.storage(), 10, size=torch.Size([2, 2]), stride=[7, 1]) + self.assertEqual( + strided[ri([0, 1]), ri([1, 0])], + torch.tensor([11, 17], dtype=dtype, device=device), + ) + strided[ri([0, 1]), ri([1, 0])] = torch.tensor( + [-1, 2], dtype=dtype, device=device + ) + self.assertEqual( + strided[ri([0, 1]), ri([1, 0])], + torch.tensor([-1, 2], dtype=dtype, device=device), + ) - reference = torch.arange(0., 24, dtype=dtype, device=device).view(3, 8) + reference = torch.arange(0.0, 24, dtype=dtype, device=device).view(3, 8) strided = torch.tensor((), dtype=dtype, device=device) - strided.set_(reference.storage(), 10, size=torch.Size([2, 2]), - stride=[7, 1]) + strided.set_(reference.storage(), 10, size=torch.Size([2, 2]), stride=[7, 1]) - rows = ri([[0], - [1]]) - columns = ri([[0, 1], - [0, 1]]) - self.assertEqual(strided[rows, columns], - torch.tensor([[10, 11], [17, 18]], dtype=dtype, device=device)) - strided[rows, columns] = torch.tensor([[4, 6], [2, 3]], dtype=dtype, device=device) - self.assertEqual(strided[rows, columns], - torch.tensor([[4, 6], [2, 3]], dtype=dtype, device=device)) + rows = ri([[0], [1]]) + columns = ri([[0, 1], [0, 1]]) + self.assertEqual( + strided[rows, columns], + torch.tensor([[10, 11], [17, 18]], dtype=dtype, device=device), + ) + strided[rows, columns] = torch.tensor( + [[4, 6], [2, 3]], dtype=dtype, device=device + ) + self.assertEqual( + strided[rows, columns], + torch.tensor([[4, 6], [2, 3]], dtype=dtype, device=device), + ) # Tests using less than the number of dims, and ellipsis @@ -401,12 +550,17 @@ class TestIndexing(TestCase): # 3 4 # 5 6 reference = consec((3, 2)) - self.assertEqual(reference[ri([0, 2]), ], - torch.tensor([[1, 2], [5, 6]], dtype=dtype, device=device)) - self.assertEqual(reference[ri([1]), ...], - torch.tensor([[3, 4]], dtype=dtype, device=device)) - self.assertEqual(reference[..., ri([1])], - torch.tensor([[2], [4], [6]], dtype=dtype, device=device)) + self.assertEqual( + reference[ri([0, 2]),], + torch.tensor([[1, 2], [5, 6]], dtype=dtype, device=device), + ) + self.assertEqual( + reference[ri([1]), ...], torch.tensor([[3, 4]], dtype=dtype, device=device) + ) + self.assertEqual( + reference[..., ri([1])], + torch.tensor([[2], [4], [6]], dtype=dtype, device=device), + ) # verify too many indices fails with self.assertRaises(IndexError): @@ -417,21 +571,22 @@ class TestIndexing(TestCase): # can't test cuda because it is a device assert if not reference.is_cuda: for err_idx in (10, -11): - with self.assertRaisesRegex(IndexError, r'out of'): + with self.assertRaisesRegex(IndexError, r"out of"): reference[err_idx] - with self.assertRaisesRegex(IndexError, r'out of'): + with self.assertRaisesRegex(IndexError, r"out of"): reference[torch.LongTensor([err_idx]).to(device)] - with self.assertRaisesRegex(IndexError, r'out of'): + with self.assertRaisesRegex(IndexError, r"out of"): reference[[err_idx]] def tensor_indices_to_np(tensor, indices): # convert the Torch Tensor to a numpy array - tensor = tensor.to(device='cpu') + tensor = tensor.to(device="cpu") npt = tensor.numpy() # convert indices - idxs = tuple(i.tolist() if isinstance(i, torch.LongTensor) else - i for i in indices) + idxs = tuple( + i.tolist() if isinstance(i, torch.LongTensor) else i for i in indices + ) return npt, idxs @@ -443,7 +598,7 @@ class TestIndexing(TestCase): def set_numpy(tensor, indices, value): if not isinstance(value, int): - if self.device_type != 'cpu': + if self.device_type != "cpu": value = value.cpu() value = value.numpy() @@ -458,7 +613,9 @@ class TestIndexing(TestCase): pyt = tensor.clone() numt = tensor.clone() pyt[indexer] = val - numt = torch.tensor(set_numpy(numt, indexer, val), dtype=dtype, device=device) + numt = torch.tensor( + set_numpy(numt, indexer, val), dtype=dtype, device=device + ) self.assertEqual(pyt, numt) def assert_backward_eq(tensor, indexer): @@ -481,18 +638,15 @@ class TestIndexing(TestCase): # 5 6 7 8 9 # 10 11 12 13 14 # 15 16 17 18 19 - reference = torch.arange(0., 20, dtype=dtype, device=device).view(4, 5) + reference = torch.arange(0.0, 20, dtype=dtype, device=device).view(4, 5) indices_to_test = [ # grab the second, fourth columns [slice(None), [1, 3]], - # first, third rows, [[0, 2], slice(None)], - # weird shape - [slice(None), [[0, 1], - [2, 3]]], + [slice(None), [[0, 1], [2, 3]]], # negatives [[-1], [0]], [[0, 2], [-1]], @@ -504,16 +658,14 @@ class TestIndexing(TestCase): for indexer in get_indices_to_test: assert_get_eq(reference, indexer) - if self.device_type != 'cpu': + if self.device_type != "cpu": assert_backward_eq(reference, indexer) for indexer in indices_to_test: assert_set_eq(reference, indexer, 44) - assert_set_eq(reference, - indexer, - get_set_tensor(reference, indexer)) + assert_set_eq(reference, indexer, get_set_tensor(reference, indexer)) - reference = torch.arange(0., 160, dtype=dtype, device=device).view(4, 8, 5) + reference = torch.arange(0.0, 160, dtype=dtype, device=device).view(4, 8, 5) indices_to_test = [ [slice(None), slice(None), [0, 3, 4]], @@ -537,7 +689,9 @@ class TestIndexing(TestCase): [[0, 2, 3], slice(None), [1, 3, 4]], # [...] # less dim, ellipsis - [[0, 2], ], + [ + [0, 2], + ], [[0, 2], slice(None)], [[0, 2], Ellipsis], [[0, 2], slice(None), Ellipsis], @@ -548,7 +702,6 @@ class TestIndexing(TestCase): [Ellipsis, [2, 3, 4]], [Ellipsis, slice(None), [2, 3, 4]], [slice(None), Ellipsis, [2, 3, 4]], - # ellipsis counts for nothing [Ellipsis, slice(None), slice(None), [0, 3, 4]], [slice(None), Ellipsis, slice(None), [0, 3, 4]], @@ -566,7 +719,7 @@ class TestIndexing(TestCase): if torch.cuda.is_available(): assert_backward_eq(reference, indexer) - reference = torch.arange(0., 1296, dtype=dtype, device=device).view(3, 9, 8, 6) + reference = torch.arange(0.0, 1296, dtype=dtype, device=device).view(3, 9, 8, 6) indices_to_test = [ [slice(None), slice(None), slice(None), [0, 3, 4]], @@ -610,7 +763,6 @@ class TestIndexing(TestCase): [[0], [4], [1, 3, 4], slice(None)], [[1], [0, 2, 3], [1], slice(None)], [[[1, 2], [1, 2]], [[0, 1], [2, 3]], [[2, 3], [3, 5]], slice(None)], - # less dim, ellipsis [Ellipsis, [0, 3, 4]], [Ellipsis, slice(None), [0, 3, 4]], @@ -624,7 +776,9 @@ class TestIndexing(TestCase): [[0], [1, 2, 4], slice(None)], [[0], [1, 2, 4], Ellipsis], [[0], [1, 2, 4], Ellipsis, slice(None)], - [[1], ], + [ + [1], + ], [[0, 2, 1], [3], [4]], [[0, 2, 1], [3], [4], slice(None)], [[0, 2, 1], [3], [4], Ellipsis], @@ -642,14 +796,16 @@ class TestIndexing(TestCase): for indexer in indices_to_test: assert_get_eq(reference, indexer) assert_set_eq(reference, indexer, 1333) - if self.device_type != 'cpu': + if self.device_type != "cpu": assert_backward_eq(reference, indexer) def test_advancedindex_big(self, device): reference = torch.arange(0, 123344, dtype=torch.int, device=device) - self.assertEqual(reference[[0, 123, 44488, 68807, 123343], ], - torch.tensor([0, 123, 44488, 68807, 123343], dtype=torch.int)) + self.assertEqual( + reference[[0, 123, 44488, 68807, 123343],], + torch.tensor([0, 123, 44488, 68807, 123343], dtype=torch.int), + ) def test_set_item_to_scalar_tensor(self, device): m = random.randint(1, 10) @@ -687,31 +843,37 @@ class TestIndexing(TestCase): def test_step_assignment(self, device): v = torch.zeros(4, 4, device=device) - v[0, 1::2] = torch.tensor([3., 4.], device=device) + v[0, 1::2] = torch.tensor([3.0, 4.0], device=device) self.assertEqual(v[0].tolist(), [0, 3, 0, 4]) self.assertEqual(v[1:].sum(), 0) def test_bool_indices(self, device): v = torch.randn(5, 7, 3, device=device) - boolIndices = torch.tensor([True, False, True, True, False], dtype=torch.bool, device=device) + boolIndices = torch.tensor( + [True, False, True, True, False], dtype=torch.bool, device=device + ) self.assertEqual(v[boolIndices].shape, (3, 7, 3)) self.assertEqual(v[boolIndices], torch.stack([v[0], v[2], v[3]])) v = torch.tensor([True, False, True], dtype=torch.bool, device=device) - boolIndices = torch.tensor([True, False, False], dtype=torch.bool, device=device) + boolIndices = torch.tensor( + [True, False, False], dtype=torch.bool, device=device + ) uint8Indices = torch.tensor([1, 0, 0], dtype=torch.uint8, device=device) with warnings.catch_warnings(record=True) as w: v1 = v[boolIndices] v2 = v[uint8Indices] self.assertEqual(v1.shape, v2.shape) self.assertEqual(v1, v2) - self.assertEqual(v[boolIndices], tensor([True], dtype=torch.bool, device=device)) + self.assertEqual( + v[boolIndices], tensor([True], dtype=torch.bool, device=device) + ) self.assertEqual(len(w), 1) def test_bool_indices_accumulate(self, device): - mask = torch.zeros(size=(10, ), dtype=torch.bool, device=device) + mask = torch.zeros(size=(10,), dtype=torch.bool, device=device) y = torch.ones(size=(10, 10), device=device) - y.index_put_((mask, ), y[mask], accumulate=True) + y.index_put_((mask,), y[mask], accumulate=True) self.assertEqual(y, torch.ones(size=(10, 10), device=device)) def test_multiple_bool_indices(self, device): @@ -730,29 +892,33 @@ class TestIndexing(TestCase): self.assertEqual(res, torch.stack([v[0], v[2], v[3]])) self.assertEqual(len(w), 1) - v = torch.tensor([1.], device=device) + v = torch.tensor([1.0], device=device) self.assertEqual(v[v == 0], torch.tensor([], device=device)) def test_byte_mask_accumulate(self, device): - mask = torch.zeros(size=(10, ), dtype=torch.uint8, device=device) + mask = torch.zeros(size=(10,), dtype=torch.uint8, device=device) y = torch.ones(size=(10, 10), device=device) with warnings.catch_warnings(record=True) as w: warnings.simplefilter("always") - y.index_put_((mask, ), y[mask], accumulate=True) + y.index_put_((mask,), y[mask], accumulate=True) self.assertEqual(y, torch.ones(size=(10, 10), device=device)) self.assertEqual(len(w), 2) - @skipIfTorchDynamo("This test causes SIGKILL when running with dynamo, https://github.com/pytorch/pytorch/issues/88472") + @skipIfTorchDynamo( + "This test causes SIGKILL when running with dynamo, https://github.com/pytorch/pytorch/issues/88472" + ) @serialTest(TEST_CUDA) def test_index_put_accumulate_large_tensor(self, device): # 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, device=device) - indices = torch.tensor([-2, 0, -2, -1, 0, -1, 1], device=device, dtype=torch.long) + indices = torch.tensor( + [-2, 0, -2, -1, 0, -1, 1], device=device, dtype=torch.long + ) values = torch.tensor([6, 5, 6, 6, 5, 7, 11], dtype=dt, device=device) - a.index_put_((indices, ), values, accumulate=True) + a.index_put_((indices,), values, accumulate=True) self.assertEqual(a[0], 11) self.assertEqual(a[1], 12) @@ -787,11 +953,19 @@ class TestIndexing(TestCase): t_dev = t.to(device) indices = [ torch.tensor([0, 1, 2, 3]), - torch.tensor([1, ]), + torch.tensor( + [ + 1, + ] + ), ] indices_dev = [i.to(device) for i in indices] values0d = torch.tensor(1.0) - values1d = torch.tensor([1.0, ]) + values1d = torch.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) @@ -805,13 +979,21 @@ class TestIndexing(TestCase): t_dev = t.to(device) indices = [ - torch.tensor([0, ]), + torch.tensor( + [ + 0, + ] + ), torch.arange(3)[:, None], torch.arange(2)[None, :], ] indices_dev = [i.to(device) for i in indices] values1d = torch.tensor([-1.0, -2.0]) - values2d = torch.tensor([[-1.0, -2.0], ]) + values2d = torch.tensor( + [ + [-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) @@ -830,7 +1012,9 @@ class TestIndexing(TestCase): self.assertTrue(not t1.is_contiguous()) self.assertTrue(not t2.is_contiguous()) - indices = [torch.tensor([0, 1]), ] + indices = [ + torch.tensor([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) @@ -952,7 +1136,11 @@ class TestIndexing(TestCase): scripted_fn2 = torch.jit.script(fn2) data = torch.arange(100, device=device, dtype=torch.float) out = scripted_fn1(data.detach().clone()) - ref = torch.tensor(np.concatenate((np.ones(50), np.arange(50, 100))), device=device, dtype=torch.float) + ref = torch.tensor( + np.concatenate((np.ones(50), np.arange(50, 100))), + device=device, + dtype=torch.float, + ) self.assertEqual(out, ref) out = scripted_fn2(data.detach().clone()) self.assertEqual(out, ref) @@ -963,9 +1151,15 @@ class TestIndexing(TestCase): self.assertEqual(v[:, [0, 4, 2]].shape, (5, 3, 3)) self.assertEqual(v[:, [[0, 1], [4, 3]]].shape, (5, 2, 2, 3)) - @dtypes(torch.cfloat, torch.cdouble, torch.float, torch.bfloat16, torch.long, torch.bool) - @dtypesIfCPU(torch.cfloat, torch.cdouble, torch.float, torch.long, torch.bool, torch.bfloat16) - @dtypesIfCUDA(torch.cfloat, torch.cdouble, torch.half, torch.long, torch.bool, torch.bfloat16) + @dtypes( + torch.cfloat, torch.cdouble, torch.float, torch.bfloat16, torch.long, torch.bool + ) + @dtypesIfCPU( + torch.cfloat, torch.cdouble, torch.float, torch.long, torch.bool, torch.bfloat16 + ) + @dtypesIfCUDA( + torch.cfloat, torch.cdouble, torch.half, torch.long, torch.bool, torch.bfloat16 + ) def test_index_put_src_datatype(self, device, dtype): src = torch.ones(3, 2, 4, device=device, dtype=dtype) vals = torch.ones(3, 2, 4, device=device, dtype=dtype) @@ -1016,21 +1210,28 @@ class TestIndexing(TestCase): def test_empty_ndim_index(self, device): x = torch.randn(5, device=device) - self.assertEqual(torch.empty(0, 2, device=device), x[torch.empty(0, 2, dtype=torch.int64, device=device)]) + self.assertEqual( + torch.empty(0, 2, device=device), + x[torch.empty(0, 2, dtype=torch.int64, device=device)], + ) x = torch.randn(2, 3, 4, 5, device=device) - self.assertEqual(torch.empty(2, 0, 6, 4, 5, device=device), - x[:, torch.empty(0, 6, dtype=torch.int64, device=device)]) + self.assertEqual( + torch.empty(2, 0, 6, 4, 5, device=device), + x[:, torch.empty(0, 6, dtype=torch.int64, device=device)], + ) x = torch.empty(10, 0, device=device) self.assertEqual(x[[1, 2]].shape, (2, 0)) self.assertEqual(x[[], []].shape, (0,)) - with self.assertRaisesRegex(IndexError, 'for dimension with size 0'): + with self.assertRaisesRegex(IndexError, "for dimension with size 0"): x[:, [0, 1]] def test_empty_ndim_index_bool(self, device): x = torch.randn(5, device=device) - self.assertRaises(IndexError, lambda: x[torch.empty(0, 2, dtype=torch.uint8, device=device)]) + self.assertRaises( + IndexError, lambda: x[torch.empty(0, 2, dtype=torch.uint8, device=device)] + ) def test_empty_slice(self, device): x = torch.randn(2, 3, 4, 5, device=device) @@ -1045,7 +1246,7 @@ class TestIndexing(TestCase): true = torch.tensor(1, dtype=torch.uint8, device=device) false = torch.tensor(0, dtype=torch.uint8, device=device) - tensors = [torch.randn(2, 3, device=device), torch.tensor(3., device=device)] + tensors = [torch.randn(2, 3, device=device), torch.tensor(3.0, device=device)] for a in tensors: self.assertNotEqual(a.data_ptr(), a[True].data_ptr()) @@ -1178,18 +1379,18 @@ class TestIndexing(TestCase): self.assertEqual(x.tolist(), [[0, 1], [5, 6]]) def test_byte_tensor_assignment(self, device): - x = torch.arange(0., 16, device=device).view(4, 4) + x = torch.arange(0.0, 16, device=device).view(4, 4) b = torch.ByteTensor([True, False, True, False]).to(device) - value = torch.tensor([3., 4., 5., 6.], device=device) + value = torch.tensor([3.0, 4.0, 5.0, 6.0], device=device) with warnings.catch_warnings(record=True) as w: x[b] = value self.assertEqual(len(w), 1) self.assertEqual(x[0], value) - self.assertEqual(x[1], torch.arange(4., 8, device=device)) + self.assertEqual(x[1], torch.arange(4.0, 8, device=device)) self.assertEqual(x[2], value) - self.assertEqual(x[3], torch.arange(12., 16, device=device)) + self.assertEqual(x[3], torch.arange(12.0, 16, device=device)) def test_variable_slicing(self, device): x = torch.arange(0, 16, device=device).view(4, 4) @@ -1200,50 +1401,62 @@ class TestIndexing(TestCase): def test_ellipsis_tensor(self, device): x = torch.arange(0, 9, device=device).view(3, 3) idx = torch.tensor([0, 2], device=device) - self.assertEqual(x[..., idx].tolist(), [[0, 2], - [3, 5], - [6, 8]]) - self.assertEqual(x[idx, ...].tolist(), [[0, 1, 2], - [6, 7, 8]]) + self.assertEqual(x[..., idx].tolist(), [[0, 2], [3, 5], [6, 8]]) + self.assertEqual(x[idx, ...].tolist(), [[0, 1, 2], [6, 7, 8]]) def test_unravel_index_errors(self, device): with self.assertRaisesRegex(TypeError, r"expected 'indices' to be integer"): - torch.unravel_index( - torch.tensor(0.5, device=device), - (2, 2)) + torch.unravel_index(torch.tensor(0.5, device=device), (2, 2)) with self.assertRaisesRegex(TypeError, r"expected 'indices' to be integer"): - torch.unravel_index( - torch.tensor([], device=device), - (10, 3, 5)) + torch.unravel_index(torch.tensor([], device=device), (10, 3, 5)) - with self.assertRaisesRegex(TypeError, r"expected 'shape' to be int or sequence"): + with self.assertRaisesRegex( + TypeError, r"expected 'shape' to be int or sequence" + ): torch.unravel_index( torch.tensor([1], device=device, dtype=torch.int64), - torch.tensor([1, 2, 3])) + torch.tensor([1, 2, 3]), + ) - with self.assertRaisesRegex(TypeError, r"expected 'shape' sequence to only contain ints"): + with self.assertRaisesRegex( + TypeError, r"expected 'shape' sequence to only contain ints" + ): torch.unravel_index( - torch.tensor([1], device=device, dtype=torch.int64), - (1, 2, 2.0)) + torch.tensor([1], device=device, dtype=torch.int64), (1, 2, 2.0) + ) - with self.assertRaisesRegex(ValueError, r"'shape' cannot have negative values, but got \(2, -3\)"): - torch.unravel_index( - torch.tensor(0, device=device), - (2, -3)) + with self.assertRaisesRegex( + ValueError, r"'shape' cannot have negative values, but got \(2, -3\)" + ): + torch.unravel_index(torch.tensor(0, device=device), (2, -3)) def test_invalid_index(self, device): x = torch.arange(0, 16, device=device).view(4, 4) - self.assertRaisesRegex(TypeError, 'slice indices', lambda: x["0":"1"]) + self.assertRaisesRegex(TypeError, "slice indices", lambda: x["0":"1"]) def test_out_of_bound_index(self, device): x = torch.arange(0, 100, device=device).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]) + 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, device): x = torch.tensor(10, device=device) @@ -1253,16 +1466,19 @@ class TestIndexing(TestCase): print(x[0]) return x[0] - self.assertRaisesRegex(IndexError, 'invalid index', runner) + self.assertRaisesRegex(IndexError, "invalid index", runner) @onlyCUDA def test_invalid_device(self, device): idx = torch.tensor([0, 1]) b = torch.zeros(5, device=device) - c = torch.tensor([1., 2.], device="cpu") + c = torch.tensor([1.0, 2.0], device="cpu") for accumulate in [True, False]: - self.assertRaises(RuntimeError, lambda: torch.index_put_(b, (idx,), c, accumulate=accumulate)) + self.assertRaises( + RuntimeError, + lambda: torch.index_put_(b, (idx,), c, accumulate=accumulate), + ) @onlyCUDA def test_cpu_indices(self, device): @@ -1287,7 +1503,9 @@ class TestIndexing(TestCase): for shape in [(3, 2), (2, 3, 5), (2, 4, 0), (2, 3, 1, 4)]: for noncontiguous in [True, False]: - t = make_tensor(shape, device=device, dtype=dtype, noncontiguous=noncontiguous) + t = make_tensor( + shape, device=device, dtype=dtype, noncontiguous=noncontiguous + ) for dim in list(range(t.ndim)) + [None]: if dim is None: indices = torch.argsort(t.view(-1)) @@ -1316,8 +1534,9 @@ class TestIndexing(TestCase): indices = torch.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"): + with self.assertRaisesRegex( + RuntimeError, "input and indices should have the same number of dimensions" + ): torch.take_along_dim(t, indices[0], dim=0) # invalid `indices` dtype @@ -1345,18 +1564,26 @@ class TestIndexing(TestCase): t = make_tensor(shape, device=device, dtype=dtype) indices = torch.argsort(t, dim=dim) - with self.assertRaisesRegex(RuntimeError, "Expected all tensors to be on the same device"): + 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()"): + 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"): + 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()"): + 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) @onlyCUDA @@ -1409,7 +1636,6 @@ class TestIndexing(TestCase): self.assertRaises(IndexError, lambda: t[idx_max]) - # The tests below are from NumPy test_indexing.py with some modifications to # make them compatible with PyTorch. It's licensed under the BDS license below: # @@ -1444,9 +1670,10 @@ class TestIndexing(TestCase): # (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + class NumpyTests(TestCase): def test_index_no_floats(self, device): - a = torch.tensor([[[5.]]], device=device) + a = torch.tensor([[[5.0]]], device=device) self.assertRaises(IndexError, lambda: a[0.0]) self.assertRaises(IndexError, lambda: a[0, 0.0]) @@ -1494,9 +1721,7 @@ class NumpyTests(TestCase): self.assertRaises(IndexError, lambda: a[b]) def test_ellipsis_index(self, device): - a = tensor([[1, 2, 3], - [4, 5, 6], - [7, 8, 9]], device=device) + a = tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]], device=device) self.assertIsNot(a[...], a) self.assertEqual(a[...], a) # `a[...]` was `a` in numpy <1.9. @@ -1519,9 +1744,7 @@ class NumpyTests(TestCase): def test_single_int_index(self, device): # Single integer index selects one row - a = tensor([[1, 2, 3], - [4, 5, 6], - [7, 8, 9]], device=device) + a = tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]], device=device) self.assertEqual(a[0], [1, 2, 3]) self.assertEqual(a[-1], [7, 8, 9]) @@ -1533,9 +1756,7 @@ class NumpyTests(TestCase): def test_single_bool_index(self, device): # Single boolean index - a = tensor([[1, 2, 3], - [4, 5, 6], - [7, 8, 9]], device=device) + a = tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]], device=device) self.assertEqual(a[True], a[None]) self.assertEqual(a[False], a[None][0:0]) @@ -1544,24 +1765,24 @@ class NumpyTests(TestCase): arr = torch.ones((5, 4, 3), device=device) index = tensor([True], device=device) - self.assertRaisesRegex(IndexError, 'mask', lambda: arr[index]) + self.assertRaisesRegex(IndexError, "mask", lambda: arr[index]) index = tensor([False] * 6, device=device) - self.assertRaisesRegex(IndexError, 'mask', lambda: arr[index]) + 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)]) + self.assertRaisesRegex(IndexError, "mask", lambda: arr[index]) + self.assertRaisesRegex(IndexError, "mask", lambda: arr[(slice(None), index)]) def test_boolean_indexing_onedim(self, device): # Indexing a 2-dimensional array with # boolean array of length one - a = tensor([[0., 0., 0.]], device=device) + a = tensor([[0.0, 0.0, 0.0]], device=device) b = tensor([True], device=device) self.assertEqual(a[b], a) # boolean assignment - a[b] = 1. - self.assertEqual(a, tensor([[1., 1., 1.]], device=device)) + a[b] = 1.0 + self.assertEqual(a, tensor([[1.0, 1.0, 1.0]], device=device)) def test_boolean_assignment_value_mismatch(self, device): # A boolean assignment should fail when the shape of the values @@ -1571,34 +1792,33 @@ class NumpyTests(TestCase): def f(a, v): a[a > -1] = tensor(v).to(device) - 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, device): # Indexing a 2-dimensional array with # 2-dimensional boolean array - a = tensor([[1, 2, 3], - [4, 5, 6], - [7, 8, 9]], device=device) - b = tensor([[True, False, True], - [False, True, False], - [True, False, True]], device=device) + a = tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]], device=device) + b = tensor( + [[True, False, True], [False, True, False], [True, False, True]], + device=device, + ) self.assertEqual(a[b], tensor([1, 3, 5, 7, 9], device=device)) self.assertEqual(a[b[1]], tensor([[4, 5, 6]], device=device)) self.assertEqual(a[b[0]], a[b[2]]) # boolean assignment a[b] = 0 - self.assertEqual(a, tensor([[0, 2, 0], - [4, 0, 6], - [0, 8, 0]], device=device)) + self.assertEqual(a, tensor([[0, 2, 0], [4, 0, 6], [0, 8, 0]], device=device)) def test_boolean_indexing_weirdness(self, device): # Weird boolean indexing things a = torch.ones((2, 3, 4), device=device) self.assertEqual((0, 2, 3, 4), a[False, True, ...].shape) - self.assertEqual(torch.ones(1, 2, device=device), a[True, [0, 1], True, True, [1], [[2]]]) + self.assertEqual( + torch.ones(1, 2, device=device), a[True, [0, 1], True, True, [1], [[2]]] + ) self.assertRaises(IndexError, lambda: a[False, [0, 1], ...]) def test_boolean_indexing_weirdness_tensors(self, device): @@ -1607,7 +1827,9 @@ class NumpyTests(TestCase): true = torch.tensor(True, device=device) a = torch.ones((2, 3, 4), device=device) self.assertEqual((0, 2, 3, 4), a[False, True, ...].shape) - self.assertEqual(torch.ones(1, 2, device=device), a[true, [0, 1], true, true, [1], [[2]]]) + self.assertEqual( + torch.ones(1, 2, device=device), a[true, [0, 1], true, true, [1], [[2]]] + ) self.assertRaises(IndexError, lambda: a[false, [0, 1], ...]) def test_boolean_indexing_alldims(self, device): @@ -1619,9 +1841,7 @@ class NumpyTests(TestCase): def test_boolean_list_indexing(self, device): # Indexing a 2-dimensional array with # boolean lists - a = tensor([[1, 2, 3], - [4, 5, 6], - [7, 8, 9]], device=device) + a = tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]], device=device) b = [True, False, False] c = [True, True, False] self.assertEqual(a[b], tensor([[1, 2, 3]], device=device)) @@ -1639,14 +1859,18 @@ class NumpyTests(TestCase): def test_broaderrors_indexing(self, device): a = torch.zeros(5, 5, device=device) - 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) + 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, device): a = torch.zeros(5, device=device) ind = torch.ones(20, dtype=torch.int64, device=device) if a.is_cuda: - raise unittest.SkipTest('CUDA asserts instead of raising an exception') + 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) @@ -1658,13 +1882,13 @@ class NumpyTests(TestCase): def test_index_is_larger(self, device): # Simple case of fancy index broadcasting of the index. a = torch.zeros((5, 5), device=device) - a[[[0], [1], [2]], [0, 1, 2]] = tensor([2., 3., 4.], device=device) + a[[[0], [1], [2]], [0, 1, 2]] = tensor([2.0, 3.0, 4.0], device=device) - self.assertTrue((a[:3, :3] == tensor([2., 3., 4.], device=device)).all()) + self.assertTrue((a[:3, :3] == tensor([2.0, 3.0, 4.0], device=device)).all()) def test_broadcast_subspace(self, device): a = torch.zeros((100, 100), device=device) - v = torch.arange(0., 100, device=device)[:, None] + v = torch.arange(0.0, 100, device=device)[:, None] b = torch.arange(99, -1, -1, device=device).long() a[b] = v expected = b.float().unsqueeze(1).expand(100, 100) @@ -1679,8 +1903,9 @@ class NumpyTests(TestCase): torch.diagonal(kernel2).copy_(torch.square(col_max.view(4))) self.assertEqual(kernel, kernel2) -instantiate_device_type_tests(TestIndexing, globals(), except_for='meta') -instantiate_device_type_tests(NumpyTests, globals(), except_for='meta') -if __name__ == '__main__': +instantiate_device_type_tests(TestIndexing, globals(), except_for="meta") +instantiate_device_type_tests(NumpyTests, globals(), except_for="meta") + +if __name__ == "__main__": run_tests() diff --git a/test/test_view_ops.py b/test/test_view_ops.py index 153b65b203d2..cb019057ae9f 100644 --- a/test/test_view_ops.py +++ b/test/test_view_ops.py @@ -1,22 +1,40 @@ # Owner(s): ["module: tests"] -import torch -import numpy as np - -import unittest -from itertools import product, permutations, combinations -from functools import partial import random +import unittest +from functools import partial +from itertools import combinations, permutations, product + +import numpy as np +import torch + from torch.testing import make_tensor -from torch.testing._internal.common_utils import ( - IS_FBCODE, TestCase, run_tests, suppress_warnings, gradcheck, gradgradcheck, - numpy_to_torch_dtype_dict, skipIfTorchDynamo +from torch.testing._internal.common_device_type import ( + dtypes, + instantiate_device_type_tests, + onlyCPU, + onlyNativeDeviceTypes, + skipLazy, + skipMeta, + skipXLA, ) -from torch.testing._internal.common_device_type import \ - (instantiate_device_type_tests, onlyCPU, dtypes, onlyNativeDeviceTypes, skipLazy, skipMeta, skipXLA) from torch.testing._internal.common_dtype import ( - all_types_and_complex_and, complex_types, all_types_and, floating_and_complex_types_and, + all_types_and, + all_types_and_complex_and, + complex_types, + floating_and_complex_types_and, ) +from torch.testing._internal.common_utils import ( + gradcheck, + gradgradcheck, + IS_FBCODE, + numpy_to_torch_dtype_dict, + run_tests, + skipIfTorchDynamo, + suppress_warnings, + TestCase, +) + # TODO: replace this with make_tensor() in common_utils.py def _generate_input(shape, dtype, device, with_extremal): @@ -29,17 +47,19 @@ def _generate_input(shape, dtype, device, with_extremal): x = torch.randn(*shape, device=device) * random.randint(30, 100) x = x.to(torch.bfloat16) else: - x = torch.randn(*shape, dtype=dtype, device=device) * random.randint(30, 100) + x = torch.randn(*shape, dtype=dtype, device=device) * random.randint( + 30, 100 + ) x[torch.randn(*shape) > 0.5] = 0 if with_extremal and dtype.is_floating_point: # Use extremal values - x[torch.randn(*shape) > 0.5] = float('nan') - x[torch.randn(*shape) > 0.5] = float('inf') - x[torch.randn(*shape) > 0.5] = float('-inf') + x[torch.randn(*shape) > 0.5] = float("nan") + x[torch.randn(*shape) > 0.5] = float("inf") + x[torch.randn(*shape) > 0.5] = float("-inf") elif with_extremal and dtype.is_complex: - x[torch.randn(*shape) > 0.5] = complex('nan') - x[torch.randn(*shape) > 0.5] = complex('inf') - x[torch.randn(*shape) > 0.5] = complex('-inf') + x[torch.randn(*shape) > 0.5] = complex("nan") + x[torch.randn(*shape) > 0.5] = complex("inf") + x[torch.randn(*shape) > 0.5] = complex("-inf") elif dtype == torch.bool: x = torch.zeros(shape, dtype=dtype, device=device) x[torch.randn(*shape) > 0.5] = True @@ -48,6 +68,7 @@ def _generate_input(shape, dtype, device, with_extremal): return x + # TODO: replace this with make_tensor() in common_utils.py def _rand_shape(dim, min_size, max_size): shape = [] @@ -55,13 +76,15 @@ def _rand_shape(dim, min_size, max_size): shape.append(random.randint(min_size, max_size)) return tuple(shape) + # TODO: refactor tests to avoid this function # Converts half/bfloat16 dtype to float when device is cpu def _convert_t(dtype, device): - if device == 'cpu' and dtype in {torch.half, torch.bfloat16}: + if device == "cpu" and dtype in {torch.half, torch.bfloat16}: return torch.float return dtype + # TODO: replace this with make_tensor() in common_utils.py # Returns a tensor of the requested shape, dtype, and device # Requesting a half CPU tensor returns a float CPU tensor with @@ -80,28 +103,31 @@ def _make_tensor(shape, dtype, device, fill_ones=False) -> torch.Tensor: return t.to(_convert_t(dtype, device)) # Populates the CPU tensor with floats representable as half/bfloat16 - if dtype == torch.half and device == 'cpu': + if dtype == torch.half and device == "cpu": return torch.randn(*shape, dtype=torch.float, device=device).half().float() - if dtype == torch.bfloat16 and device == 'cpu': + if dtype == torch.bfloat16 and device == "cpu": return torch.randn(*shape, dtype=torch.float, device=device).bfloat16().float() # Default: returns a tensor with random float values return torch.randn(shape, dtype=dtype, device=device).to(dtype=dtype) + # Tests ops and indexing to ensure they return views (and new tensors) as # appropriate. class TestViewOps(TestCase): exact_dtype = True def is_view_of(self, base, other): - if (not other._is_view() or - other is base or - other._base is not base or - base.device != other.device): + if ( + not other._is_view() + or other is base + or other._base is not base + or base.device != other.device + ): return False # Note: only validates storage on native device types # because some accelerators, like XLA, do not expose storage - if base.device.type == 'cpu' or base.device.type == 'cuda': + if base.device.type == "cpu" or base.device.type == "cuda": if base.untyped_storage().data_ptr() != other.untyped_storage().data_ptr(): return False @@ -109,7 +135,7 @@ class TestViewOps(TestCase): # Returns true if v1 and v2 are views of the same base def is_view_of_same_base(self, v1, v2): - if (not v1._is_view() or v1 is v2): + if not v1._is_view() or v1 is v2: return False return self.is_view_of(v1._base, v2) @@ -130,15 +156,23 @@ class TestViewOps(TestCase): @onlyNativeDeviceTypes @dtypes(*all_types_and_complex_and(torch.half, torch.bool)) def test_view_dtype_new(self, device, dtype): - dtypes = {value : key for (key, value) in numpy_to_torch_dtype_dict.items()} + dtypes = {value: key for (key, value) in numpy_to_torch_dtype_dict.items()} del dtypes[torch.bool] def generate_inputs(): yield make_tensor((4, 4, 64), dtype=dtype, device=device, low=-5, high=5) - yield make_tensor((4, 4, 64), dtype=dtype, device=device, low=-5, high=5).permute(1, 0, 2) - yield make_tensor((4, 64, 4), dtype=dtype, device=device, low=-5, high=5).permute(2, 0, 1) - yield make_tensor((1, 5, 1), dtype=dtype, device=device, low=-5, high=5).expand(5, 5, 64) - yield make_tensor((2, 5, 256), dtype=dtype, device=device, low=-5, high=5)[1::2, 1:, ::2] + yield make_tensor( + (4, 4, 64), dtype=dtype, device=device, low=-5, high=5 + ).permute(1, 0, 2) + yield make_tensor( + (4, 64, 4), dtype=dtype, device=device, low=-5, high=5 + ).permute(2, 0, 1) + yield make_tensor( + (1, 5, 1), dtype=dtype, device=device, low=-5, high=5 + ).expand(5, 5, 64) + yield make_tensor((2, 5, 256), dtype=dtype, device=device, low=-5, high=5)[ + 1::2, 1:, ::2 + ] yield make_tensor((0, 5, 64), dtype=dtype, device=device, low=-5, high=5) yield make_tensor((), dtype=dtype, device=device, low=-5, high=5) @@ -174,15 +208,21 @@ class TestViewOps(TestCase): a_np_contiguous = a.cpu().contiguous().numpy() for view_dtype, np_view_dtype in dtypes.items(): - equal_element_size = torch._utils._element_size(dtype) == torch._utils._element_size(view_dtype) + equal_element_size = torch._utils._element_size( + dtype + ) == torch._utils._element_size(view_dtype) if not equal_element_size and a.dim() == 0: - with self.assertRaisesRegex(RuntimeError, r"self.dim\(\) cannot be 0"): + with self.assertRaisesRegex( + RuntimeError, r"self.dim\(\) cannot be 0" + ): a.view(view_dtype) continue if not equal_element_size and a.stride(-1) != 1: - with self.assertRaisesRegex(RuntimeError, r"self.stride\(-1\) must be 1"): + with self.assertRaisesRegex( + RuntimeError, r"self.stride\(-1\) must be 1" + ): a.view(view_dtype) continue @@ -190,7 +230,9 @@ class TestViewOps(TestCase): self.assertEqual(a_view.dtype, view_dtype) self.assertEqual(a.data_ptr(), a_view.data_ptr()) - expected_size, expected_stride = calc_expected_size_and_stride(a, view_dtype) + expected_size, expected_stride = calc_expected_size_and_stride( + a, view_dtype + ) self.assertEqual(a_view.size(), expected_size) self.assertEqual(a_view.stride(), expected_stride) @@ -210,8 +252,17 @@ class TestViewOps(TestCase): # because view(dtype) does not support backward yet # TODO: Remove this when autograd support is added if dtype.is_floating_point or dtype.is_complex: - for view_dtype in floating_and_complex_types_and(torch.half, torch.bfloat16): - t = make_tensor((5, 5, 64), dtype=dtype, device=device, low=-5, high=5, requires_grad=True) + for view_dtype in floating_and_complex_types_and( + torch.half, torch.bfloat16 + ): + t = make_tensor( + (5, 5, 64), + dtype=dtype, + device=device, + low=-5, + high=5, + requires_grad=True, + ) self.assertFalse(t.view(view_dtype).requires_grad) # Test the extra error checks that happen when the view dtype @@ -221,28 +272,35 @@ class TestViewOps(TestCase): def test_view_dtype_upsize_errors(self, device, dtype): dtype_size = torch._utils._element_size(dtype) - for view_dtype in all_types_and_complex_and(torch.half, torch.bfloat16, torch.bool): + for view_dtype in all_types_and_complex_and( + torch.half, torch.bfloat16, torch.bool + ): view_dtype_size = torch._utils._element_size(view_dtype) if view_dtype_size <= dtype_size: continue size_ratio = view_dtype_size // dtype_size - a = make_tensor((4, 4, size_ratio + 1), dtype=dtype, device=device, low=-5, high=5) + a = make_tensor( + (4, 4, size_ratio + 1), dtype=dtype, device=device, low=-5, high=5 + ) with self.assertRaisesRegex( - RuntimeError, - rf"self.size\(-1\) must be divisible by {size_ratio}"): + RuntimeError, rf"self.size\(-1\) must be divisible by {size_ratio}" + ): a.view(view_dtype) with self.assertRaisesRegex( - RuntimeError, - rf"self.storage_offset\(\) must be divisible by {size_ratio}"): + RuntimeError, + rf"self.storage_offset\(\) must be divisible by {size_ratio}", + ): a[:, :, 1:].view(view_dtype) - a = make_tensor((4, 4, size_ratio), dtype=dtype, device=device, low=-5, high=5) + a = make_tensor( + (4, 4, size_ratio), dtype=dtype, device=device, low=-5, high=5 + ) a = a.as_strided((4, 4, size_ratio), (size_ratio, 1, 1)) with self.assertRaisesRegex( - RuntimeError, - rf"self.stride\(1\) must be divisible by {size_ratio}"): + RuntimeError, rf"self.stride\(1\) must be divisible by {size_ratio}" + ): a.view(view_dtype) @onlyNativeDeviceTypes @@ -255,14 +313,18 @@ class TestViewOps(TestCase): if input.size()[-1] != 2: self.assertRaisesRegex( - RuntimeError, "Tensor must have a last dimension of size 2", - lambda: torch.view_as_complex(input)) + RuntimeError, + "Tensor must have a last dimension of size 2", + lambda: torch.view_as_complex(input), + ) return if input.stride()[-1] != 1: self.assertRaisesRegex( - RuntimeError, "Tensor must have a last dimension with stride 1", - lambda: torch.view_as_complex(input)) + RuntimeError, + "Tensor must have a last dimension with stride 1", + lambda: torch.view_as_complex(input), + ) return res = torch.view_as_complex(input) @@ -276,25 +338,30 @@ class TestViewOps(TestCase): # RuntimeError since in this case the last dim of input would not have stride 1 fn(contiguous_input=False, dim0=1, dim1=2) - # RuntimeError since in this case the stride of non-last dim of input would not be of size 2 x = torch.randn(3, 3, device=device) t = torch.as_strided(x, (2, 2), (1, 1)) self.assertRaisesRegex( - RuntimeError, "Tensor must have a stride divisible by 2 for all but last dimension", - lambda: torch.view_as_complex(t)) + RuntimeError, + "Tensor must have a stride divisible by 2 for all but last dimension", + lambda: torch.view_as_complex(t), + ) # tensor with zero elements x = torch.tensor([], device=device) # torch.Size([0]) self.assertRaisesRegex( - RuntimeError, "Tensor must have a last dimension of size 2", - lambda: torch.view_as_complex(x)) + RuntimeError, + "Tensor must have a last dimension of size 2", + lambda: torch.view_as_complex(x), + ) # zero dimension tensor z = torch.tensor(2.0) self.assertRaisesRegex( - RuntimeError, "Input tensor must have one or more dimensions", - lambda: torch.view_as_complex(z)) + RuntimeError, + "Input tensor must have one or more dimensions", + lambda: torch.view_as_complex(z), + ) y = x.reshape(0, 2) # torch.Size([0, 2]) res = torch.view_as_complex(y) @@ -410,13 +477,20 @@ class TestViewOps(TestCase): @onlyNativeDeviceTypes @dtypes(*complex_types()) def test_conj_imag_view(self, device, dtype) -> None: - t = _make_tensor((4, 5,), dtype, device) + t = _make_tensor( + ( + 4, + 5, + ), + dtype, + device, + ) t_numpy_conj = torch.from_numpy(t.cpu().numpy().conj()).to(device=device) v = t.conj() self.assertTrue(self.is_view_of(t, v)) self.assertEqual(v, t_numpy_conj) - if (t.is_complex()): + if t.is_complex(): v_imag = v.imag self.assertTrue(self.is_view_of(t, v_imag)) self.assertEqual(v_imag, t_numpy_conj.imag) @@ -424,7 +498,14 @@ class TestViewOps(TestCase): @onlyNativeDeviceTypes def test_conj_view_with_shared_memory(self, device) -> None: - a = _make_tensor((4, 5,), torch.cfloat, device) + a = _make_tensor( + ( + 4, + 5, + ), + torch.cfloat, + device, + ) b = a.conj() c = a.conj() @@ -433,7 +514,12 @@ class TestViewOps(TestCase): self.assertEqual(torch.add(b, c), b.add_(c)) @onlyNativeDeviceTypes - @dtypes(*product(complex_types(), all_types_and_complex_and(torch.half, torch.bfloat16, torch.bool))) + @dtypes( + *product( + complex_types(), + all_types_and_complex_and(torch.half, torch.bfloat16, torch.bool), + ) + ) @suppress_warnings def test_set_real_imag(self, device, dtypes): x = torch.randn(10, dtype=dtypes[0], device=device) @@ -499,9 +585,10 @@ class TestViewOps(TestCase): stacked = torch.randn(3, 10, 10, requires_grad=True) outs = stacked.unbind() gi = grad.unbind()[i] - g, = torch.autograd.grad(outs[i], stacked, gi) - g_expected = torch.stack([gi if j == i else torch.zeros_like(gi) - for j in range(3)], dim=0) + (g,) = torch.autograd.grad(outs[i], stacked, gi) + g_expected = torch.stack( + [gi if j == i else torch.zeros_like(gi) for j in range(3)], dim=0 + ) self.assertEqual(g, g_expected) # Check with gradcheck stacked = torch.randn(3, 10, 10, dtype=torch.double, requires_grad=True) @@ -789,8 +876,9 @@ class TestViewOps(TestCase): self.assertTrue(self.is_view_of_same_base(t, v)) # stride[i] = stride[i + 1] * size[i + 1] is satisfied for 3 groups: - t = torch.ones(720, device=device) \ - .as_strided((2, 3, 2, 3, 5, 4), (6, 2, 15, 5, 1, 0)) + t = torch.ones(720, device=device).as_strided( + (2, 3, 2, 3, 5, 4), (6, 2, 15, 5, 1, 0) + ) # [--1--|---2---|-3-] [--1--|----2---|-3-] v1 = t.flatten(0, 1) v2 = v1.flatten(1, 3) @@ -811,6 +899,7 @@ class TestViewOps(TestCase): nv[idx_nv] = 0 if device != "meta": self.assertNotEqual(t[idx_t], nv[idx_nv]) + t = torch.ones(2, 3, 2, 3, device=device).transpose(2, 3) nv = t.flatten(1, 3) assert_is_nonview(t, nv) @@ -858,7 +947,9 @@ class TestViewOps(TestCase): nv[1, 1] = 0 self.assertNotEqual(t[2, 2], nv[1, 1]) - @unittest.skipIf(IS_FBCODE, "TorchScript backend not yet supported in FBCODE/OVRSOURCE builds") + @unittest.skipIf( + IS_FBCODE, "TorchScript backend not yet supported in FBCODE/OVRSOURCE builds" + ) def test_advanced_indexing_assignment(self, device): t = torch.ones(3, 3, device=device) rows = torch.tensor([[0, 0], [2, 2]], device=device) @@ -953,9 +1044,9 @@ class TestViewOps(TestCase): self.assertEqual(expected1, out1) self.assertEqual(expected2, out2) + class TestOldViewOps(TestCase): def test_ravel(self, device): - def _test_ravel(tensors, size, nc=False): for src in tensors: # Continuous Tensor -> View @@ -996,34 +1087,32 @@ class TestOldViewOps(TestCase): self.assertTrue(flat2.is_contiguous()) # Test both float tensor and quantized tensor - tensors = [torch.randn(5, 5, 5, 5, device=device), - torch._empty_affine_quantized([5, 5, 5, 5], - scale=2, - zero_point=3, - dtype=torch.quint8, - device=device)] + tensors = [ + torch.randn(5, 5, 5, 5, device=device), + torch._empty_affine_quantized( + [5, 5, 5, 5], scale=2, zero_point=3, dtype=torch.quint8, device=device + ), + ] _test_ravel(tensors, 625) - tensors = [torch.randn(0, 2, 3, device=device), - torch.randn(3, 0, 2, device=device), - torch._empty_affine_quantized([0, 2, 3], - scale=2, - zero_point=3, - dtype=torch.quint8, - device=device), - torch._empty_affine_quantized([3, 0, 2], - scale=2, - zero_point=3, - dtype=torch.quint8, - device=device)] + tensors = [ + torch.randn(0, 2, 3, device=device), + torch.randn(3, 0, 2, device=device), + torch._empty_affine_quantized( + [0, 2, 3], scale=2, zero_point=3, dtype=torch.quint8, device=device + ), + torch._empty_affine_quantized( + [3, 0, 2], scale=2, zero_point=3, dtype=torch.quint8, device=device + ), + ] _test_ravel(tensors, 0) - tensors = [torch.randn(5, 5, device=device), - torch._empty_affine_quantized([5, 5], - scale=2, - zero_point=3, - dtype=torch.quint8, - device=device)] + tensors = [ + torch.randn(5, 5, device=device), + torch._empty_affine_quantized( + [5, 5], scale=2, zero_point=3, dtype=torch.quint8, device=device + ), + ] _test_ravel(tensors, 25, True) # TODO: this should be refactored into the view ops test suite @@ -1055,7 +1144,9 @@ class TestOldViewOps(TestCase): # test non-contiguous noncontig = torch.randn(5, 2, 1, 3, device=device)[:, 0] self.assertFalse(noncontig.is_contiguous()) - self.assertEqual(noncontig.expand(2, 5, 4, 3), noncontig.contiguous().repeat(2, 1, 4, 1)) + self.assertEqual( + noncontig.expand(2, 5, 4, 3), noncontig.contiguous().repeat(2, 1, 4, 1) + ) # make sure it's compatible with unsqueeze expanded = tensor2.expand(1, 1, 5) @@ -1068,7 +1159,9 @@ class TestOldViewOps(TestCase): self.assertRaises(RuntimeError, lambda: tensor2.expand(-1, -1)) # test expanding empty to empty - self.assertEqual(torch.zeros(0, device=device).expand((0,)), torch.zeros(0, device=device)) + self.assertEqual( + torch.zeros(0, device=device).expand((0,)), torch.zeros(0, device=device) + ) # TODO: this should be refactored into the view ops test suite def test_view_empty(self, device): @@ -1107,7 +1200,9 @@ class TestOldViewOps(TestCase): x = torch.randn(3, 3, device=device) self.assertEqual(x.data_ptr(), x.reshape_as(torch.rand(9)).data_ptr()) self.assertEqual(x.data_ptr(), x.reshape_as(torch.rand(1, 9, 1)).data_ptr()) - self.assertRaises(RuntimeError, lambda: x.reshape_as(torch.rand(10, device=device))) + self.assertRaises( + RuntimeError, lambda: x.reshape_as(torch.rand(10, device=device)) + ) def test_flatten(self, device): # Test that flatten returns 1-dim tensor when given a 0-dim tensor @@ -1125,12 +1220,12 @@ class TestOldViewOps(TestCase): self.assertEqual(flat0.shape, flat1.shape) # Test both float tensor and quantized tensor - tensors = [torch.randn(5, 5, 5, 5, device=device), - torch._empty_affine_quantized([5, 5, 5, 5], - scale=2, - zero_point=3, - dtype=torch.quint8, - device=device)] + tensors = [ + torch.randn(5, 5, 5, 5, device=device), + torch._empty_affine_quantized( + [5, 5, 5, 5], scale=2, zero_point=3, dtype=torch.quint8, device=device + ), + ] for src in tensors: flat = src.flatten(0, -1) self.assertEqual(flat.shape, torch.Size([625])) @@ -1160,11 +1255,13 @@ class TestOldViewOps(TestCase): self.assertEqual(flat, src) # out of bounds index - with self.assertRaisesRegex(IndexError, 'Dimension out of range'): + with self.assertRaisesRegex(IndexError, "Dimension out of range"): src.flatten(5, 10) # invalid start and end - with self.assertRaisesRegex(RuntimeError, 'start_dim cannot come after end_dim'): + with self.assertRaisesRegex( + RuntimeError, "start_dim cannot come after end_dim" + ): src.flatten(2, 0) # TODO: update to work on CUDA, too @@ -1176,7 +1273,9 @@ class TestOldViewOps(TestCase): self.assertEqual(x.narrow(0, 1, 1), torch.tensor([[3, 4, 5]])) self.assertEqual(x.narrow(0, -1, 1), torch.tensor([[6, 7, 8]])) self.assertEqual(x.narrow(0, -2, 2), torch.tensor([[3, 4, 5], [6, 7, 8]])) - self.assertEqual(x.narrow(0, -3, 3), torch.tensor([[0, 1, 2], [3, 4, 5], [6, 7, 8]])) + self.assertEqual( + x.narrow(0, -3, 3), torch.tensor([[0, 1, 2], [3, 4, 5], [6, 7, 8]]) + ) self.assertEqual(x.narrow(-1, -1, 1), torch.tensor([[2], [5], [8]])) self.assertEqual(x.narrow(-2, -1, 1), torch.tensor([[6, 7, 8]])) @@ -1186,7 +1285,7 @@ class TestOldViewOps(TestCase): x = torch.tensor([[0, 1, 2], [3, 4, 5], [6, 7, 8]]) self.assertEqual(x.narrow(0, torch.tensor(0), 1), torch.tensor([[0, 1, 2]])) with self.assertRaises(Exception): - x.narrow(0, torch.tensor(0.), 1) + x.narrow(0, torch.tensor(0.0), 1) with self.assertRaises(Exception): x.narrow(0, torch.tensor([0]), 1) with self.assertRaises(Exception): @@ -1215,10 +1314,14 @@ class TestOldViewOps(TestCase): # Test 3D tensor x = torch.rand((2, 2, 2)) - with self.assertRaisesRegex(RuntimeError, 'expects a tensor with <= 2 dimensions, but self is 3D'): + with self.assertRaisesRegex( + RuntimeError, "expects a tensor with <= 2 dimensions, but self is 3D" + ): x.t() x = x.to_sparse() - with self.assertRaisesRegex(RuntimeError, 'expects a tensor with <= 2 sparse and 0 dense dimensions'): + with self.assertRaisesRegex( + RuntimeError, "expects a tensor with <= 2 sparse and 0 dense dimensions" + ): x.t() @onlyCPU @@ -1231,19 +1334,23 @@ class TestOldViewOps(TestCase): start = 0 for target_size, split in zip(target_sizes, splits): self.assertEqual(split.size(), target_size) - self.assertEqual(tensor.narrow(dim, start, target_size[dim]), split, atol=0, rtol=0) + self.assertEqual( + tensor.narrow(dim, start, target_size[dim]), split, atol=0, rtol=0 + ) start = start + target_size[dim] # Variable sections split tensor = torch.randn(20, 10) dim = 0 split_sizes = [5, 5, 10] - target_sizes = ([[5, 10], [5, 10], [10, 10]]) + target_sizes = [[5, 10], [5, 10], [10, 10]] splits = tensor.split(split_sizes, dim) start = 0 for target_size, split in zip(target_sizes, splits): self.assertEqual(split.size(), target_size) - self.assertEqual(tensor.narrow(dim, start, target_size[dim]), split, atol=0, rtol=0) + self.assertEqual( + tensor.narrow(dim, start, target_size[dim]), split, atol=0, rtol=0 + ) start = start + target_size[dim] split_sizes = [2, 2, 6] @@ -1253,7 +1360,9 @@ class TestOldViewOps(TestCase): start = 0 for target_size, split in zip(target_sizes, splits): self.assertEqual(split.size(), target_size) - self.assertEqual(tensor.narrow(dim, start, target_size[dim]), split, atol=0, rtol=0) + self.assertEqual( + tensor.narrow(dim, start, target_size[dim]), split, atol=0, rtol=0 + ) start = start + target_size[dim] @onlyCPU @@ -1266,12 +1375,13 @@ class TestOldViewOps(TestCase): start = 0 for target_size, split in zip(target_sizes, splits): self.assertEqual(split.size(), target_size) - self.assertEqual(tensor.narrow(dim, start, target_size[dim]), split, - atol=0, rtol=0) + self.assertEqual( + tensor.narrow(dim, start, target_size[dim]), split, atol=0, rtol=0 + ) start = start + target_size[dim] # Invalid chunk sizes - error_regex = 'chunk expects.*greater than 0' + error_regex = "chunk expects.*greater than 0" with self.assertRaisesRegex(RuntimeError, error_regex): tensor.chunk(0) with self.assertRaisesRegex(RuntimeError, error_regex): @@ -1312,7 +1422,9 @@ class TestOldViewOps(TestCase): @dtypes(*all_types_and_complex_and(torch.half, torch.bfloat16, torch.bool)) def test_transposes(self, device, dtype): for op in ("T", "H", "mT", "mH", "adjoint"): - shapes = ((2, 3), (2, 3, 4)) if op[0] == "m" or op == "adjoint" else ((2, 3),) + shapes = ( + ((2, 3), (2, 3, 4)) if op[0] == "m" or op == "adjoint" else ((2, 3),) + ) for shape in shapes: a = make_tensor(shape, device=device, dtype=dtype) t1 = getattr(a, op) @@ -1357,7 +1469,9 @@ class TestOldViewOps(TestCase): def test_memory_format_resize_as(self, device): def test_helper(shape, memory_format, device): - xc = torch.randn(shape, device=device).contiguous(memory_format=memory_format) + xc = torch.randn(shape, device=device).contiguous( + memory_format=memory_format + ) flat = torch.randn(xc.numel(), device=device) flat.resize_as_(xc, memory_format=torch.preserve_format) self.assertTrue(flat.is_contiguous(memory_format=memory_format)) @@ -1372,7 +1486,9 @@ class TestOldViewOps(TestCase): self.assertTrue(flat.is_contiguous(memory_format=memory_format)) test_helper((10, 3, 32, 32), 10 * 3 * 32 * 32, torch.channels_last, device) - test_helper((3, 10, 3, 32, 32), 3 * 10 * 3 * 32 * 32, torch.channels_last_3d, device) + test_helper( + (3, 10, 3, 32, 32), 3 * 10 * 3 * 32 * 32, torch.channels_last_3d, device + ) @onlyNativeDeviceTypes @dtypes(torch.int64, torch.float, torch.complex128) @@ -1407,14 +1523,22 @@ class TestOldViewOps(TestCase): dst_dim = dst_dim - nd partial_map = { - torch.swapdims: partial(torch.swapdims, dim0=src_dim, dim1=dst_dim), - torch.swapaxes: partial(torch.swapaxes, axis0=src_dim, axis1=dst_dim), - torch.transpose: partial(torch.transpose, dim0=src_dim, dim1=dst_dim), + torch.swapdims: partial( + torch.swapdims, dim0=src_dim, dim1=dst_dim + ), + torch.swapaxes: partial( + torch.swapaxes, axis0=src_dim, axis1=dst_dim + ), + torch.transpose: partial( + torch.transpose, dim0=src_dim, dim1=dst_dim + ), } torch_fn = partial_map[fn] np_fn = partial(np.swapaxes, axis1=src_dim, axis2=dst_dim) - self.compare_with_numpy(torch_fn, np_fn, x, device=None, dtype=None) + self.compare_with_numpy( + torch_fn, np_fn, x, device=None, dtype=None + ) # Move dim to same position x = torch.randn(2, 3, 5, 7, 11) @@ -1437,11 +1561,15 @@ class TestOldViewOps(TestCase): x = _generate_input(shape, dtype, device, with_extremal) if contiguous: x = x.T - self.compare_with_numpy(torch_fn, np_fn, x, device=None, dtype=None) + self.compare_with_numpy( + torch_fn, np_fn, x, device=None, dtype=None + ) # Compare sequence input torch_sequence_x = (x,) * random.randint(3, 10) - np_sequence_x = tuple(np.array(x.detach().cpu().numpy()) for x in torch_sequence_x) + np_sequence_x = tuple( + np.array(x.detach().cpu().numpy()) for x in torch_sequence_x + ) torch_res = torch_fn(*torch_sequence_x) np_res = np_fn(*np_sequence_x) @@ -1484,7 +1612,6 @@ class TestOldViewOps(TestCase): self._test_atleast(device, torch.atleast_2d) self._test_atleast(device, torch.atleast_3d) - @onlyCPU @dtypes(torch.float) def test_broadcast_tensors(self, device, dtype): @@ -1498,7 +1625,6 @@ class TestOldViewOps(TestCase): self.assertTrue(y1.size() == expected_size) self.assertTrue(y2.size() == expected_size) - @onlyCPU def test_broadcast_shapes(self, device): examples = [(), (1,), (2,), (1, 1), (3, 1), (3, 2), (4, 1, 1), (4, 3, 2)] @@ -1520,30 +1646,48 @@ class TestOldViewOps(TestCase): res2 = torch.broadcast_tensors(*map(torch.empty, integral_inputs))[0].shape self.assertEqual(res1, res2) - inputs_with_neg_vals = [[1, 1, -12], [-1, 1], [-11, ]] + inputs_with_neg_vals = [ + [1, 1, -12], + [-1, 1], + [ + -11, + ], + ] for integral_inputs_with_neg_vals in inputs_with_neg_vals: - with self.assertRaisesRegex(RuntimeError, "Trying to create tensor with negative dimension"): + with self.assertRaisesRegex( + RuntimeError, "Trying to create tensor with negative dimension" + ): torch.broadcast_shapes(*integral_inputs_with_neg_vals) integral_inputs_error_case = [(3, 5), (2, 4, 1)] for error_input in integral_inputs_error_case: - with self.assertRaisesRegex(RuntimeError, "Shape mismatch: objects cannot be broadcast to a single shape"): + with self.assertRaisesRegex( + RuntimeError, + "Shape mismatch: objects cannot be broadcast to a single shape", + ): torch.broadcast_shapes(*error_input) negative_inputs = [(-1,), (1, -12), (4, -11), (-4, 1), (1, 1, -2)] for s0 in negative_inputs: - with self.assertRaisesRegex(RuntimeError, "Trying to create tensor with negative dimension"): + with self.assertRaisesRegex( + RuntimeError, "Trying to create tensor with negative dimension" + ): torch.broadcast_shapes(s0) for s1 in negative_inputs: - with self.assertRaisesRegex(RuntimeError, "Trying to create tensor with negative dimension"): + with self.assertRaisesRegex( + RuntimeError, "Trying to create tensor with negative dimension" + ): torch.broadcast_shapes(s0, s1) float_inputs_error_case = [(1.1, 2.0), (1.1, 1.0)] for error_case in float_inputs_error_case: for float_input in error_case: - with self.assertRaisesRegex(RuntimeError, "Input shapes " - "should be of type ints, a tuple of ints, or a list of ints"): + with self.assertRaisesRegex( + RuntimeError, + "Input shapes " + "should be of type ints, a tuple of ints, or a list of ints", + ): torch.broadcast_shapes(float_input) diff_input_types = [(1, (5,)), (3, (1,)), (1, (3, 4))] @@ -1564,9 +1708,7 @@ class TestOldViewOps(TestCase): return False return True - sizes = ( - (), (1,), (2,), (1, 1), (3, 1), (3, 2), (4, 1, 1), (4, 3, 2) - ) + sizes = ((), (1,), (2,), (1, 1), (3, 1), (3, 2), (4, 1, 1), (4, 3, 2)) for s0, s1 in combinations(sizes, r=2): t = make_tensor(s0, dtype=dtype, device=device, low=-9, high=9) t_np = t.cpu().numpy() @@ -1576,9 +1718,11 @@ class TestOldViewOps(TestCase): np_res = np.broadcast_to(t_np, s1) self.assertEqual(res, np_res) else: - with self.assertRaisesRegex(RuntimeError, - r"The expanded size of the tensor \(\d\) " - r"must match the existing size \(\d\)"): + with self.assertRaisesRegex( + RuntimeError, + r"The expanded size of the tensor \(\d\) " + r"must match the existing size \(\d\)", + ): torch.broadcast_to(t, s1) def test_view(self, device): @@ -1602,10 +1746,14 @@ class TestOldViewOps(TestCase): self.assertEqual(empty.view(-1).size(), torch.Size([0])) self.assertEqual(empty.view(10, 3, -1).size(), torch.Size([10, 3, 0])) - with self.assertRaisesRegex(RuntimeError, r"because the unspecified dimension size -1 can be any value"): + with self.assertRaisesRegex( + RuntimeError, r"because the unspecified dimension size -1 can be any value" + ): empty.view(-1, 0) - with self.assertRaisesRegex(RuntimeError, r"because the unspecified dimension size -1 can be any value"): + with self.assertRaisesRegex( + RuntimeError, r"because the unspecified dimension size -1 can be any value" + ): empty.view(3, 0, -1, 0) self.assertRaises(RuntimeError, lambda: tensor.view(15, 0)) @@ -1614,7 +1762,11 @@ class TestOldViewOps(TestCase): # test view when tensor is not contiguous in every dimension, but only # contiguous dimensions are touched. - tensor = torch.rand(4, 2, 5, 1, 6, 2, 9, 3, device=device).transpose(-1, 2).transpose(-2, 3) + tensor = ( + torch.rand(4, 2, 5, 1, 6, 2, 9, 3, device=device) + .transpose(-1, 2) + .transpose(-2, 3) + ) # size: [ 4, 2, 3, 9, 6, 2, 1, 5] # stride: [3840, 1620, 1, 3, 54, 27, 324, 324] # contiguous dim chunks: [__________, ____, ____, __________, ____, ____] @@ -1648,7 +1800,9 @@ class TestOldViewOps(TestCase): self.assertRaises(RuntimeError, lambda: tensor.view(8, 3, 54, 2, 1, 5)) # view with stride 0 dims - tensor = torch.empty(1, 1, device=device).expand(3, 4) # all dims are contiguous + tensor = torch.empty(1, 1, device=device).expand( + 3, 4 + ) # all dims are contiguous contig_tensor = tensor.clone() self.assertEqual(tensor.view(-1), contig_tensor.view(-1)) self.assertEqual(tensor.view(1, -1, 1), contig_tensor.view(1, -1, 1)) @@ -1670,7 +1824,9 @@ class TestOldViewOps(TestCase): # the copy). copy_tensor = tensor.transpose(0, 1).reshape(target) self.assertEqual(copy_tensor.size(), target) - self.assertNotEqual(tensor.storage().data_ptr(), copy_tensor.storage().data_ptr()) + self.assertNotEqual( + tensor.storage().data_ptr(), copy_tensor.storage().data_ptr() + ) def test_contiguous(self, device): x = torch.randn(1, 16, 5, 5, device=device) @@ -1700,9 +1856,11 @@ class TestOldViewOps(TestCase): a_n = a.cpu().numpy() for dim in range(-a.dim(), a.dim()): for sections in range(1, 2 * a.size(dim)): - msg = f'input_size {input_size}, sections {sections}, dim {dim}' + msg = f"input_size {input_size}, sections {sections}, dim {dim}" result1 = torch.tensor_split(a, sections, dim) - result2 = torch.tensor_split(a, torch.tensor(sections, dtype=torch.int64), dim) + result2 = torch.tensor_split( + a, torch.tensor(sections, dtype=torch.int64), dim + ) for r1, r2 in zip(result1, result2): self.assertEqual(r1.device, torch.device(device), msg=msg) self.assertEqual(r1.dtype, dtype, msg=msg) @@ -1744,9 +1902,11 @@ class TestOldViewOps(TestCase): for dim in range(-a.dim(), a.dim()): for indices in indices_args: result_1 = torch.tensor_split(a, indices, dim) - result_2 = torch.tensor_split(a, torch.tensor(indices, dtype=torch.int64), dim) + result_2 = torch.tensor_split( + a, torch.tensor(indices, dtype=torch.int64), dim + ) - msg = f'input_size {input_size}, indices {indices}, dim {dim}' + msg = f"input_size {input_size}, indices {indices}, dim {dim}" for r1, r2 in zip(result_1, result_2): self.assertEqual(r1.device, torch.device(device), msg=msg) self.assertEqual(r1.dtype, dtype, msg=msg) @@ -1762,18 +1922,46 @@ class TestOldViewOps(TestCase): S = 10 test_cases = [ # input size, sections or indices, dim, error type, error message, numpy error type - [(S,), 10, 1, IndexError, r'Dimension out of range', IndexError], - [(), 10, 0, RuntimeError, r'tensor_split expected at least a 1-dimensional tensor, ' - + 'but got a tensor with 0 dims', IndexError], - [(S,), (10,), 1, IndexError, r'Dimension out of range', IndexError], - [(), (10,), 0, RuntimeError, r'tensor_split expected at least a 1-dimensional tensor, ' - + 'but got a tensor with 0 dims', IndexError], - [(S,), 0, 0, RuntimeError, r'number of sections must be larger than 0, got 0', ValueError], - [(S,), -1, 0, RuntimeError, r'number of sections must be larger than 0, got -1', ValueError], + [(S,), 10, 1, IndexError, r"Dimension out of range", IndexError], + [ + (), + 10, + 0, + RuntimeError, + r"tensor_split expected at least a 1-dimensional tensor, " + + "but got a tensor with 0 dims", + IndexError, + ], + [(S,), (10,), 1, IndexError, r"Dimension out of range", IndexError], + [ + (), + (10,), + 0, + RuntimeError, + r"tensor_split expected at least a 1-dimensional tensor, " + + "but got a tensor with 0 dims", + IndexError, + ], + [ + (S,), + 0, + 0, + RuntimeError, + r"number of sections must be larger than 0, got 0", + ValueError, + ], + [ + (S,), + -1, + 0, + RuntimeError, + r"number of sections must be larger than 0, got -1", + ValueError, + ], ] for input_size, sections_or_indices, dim, err, err_msg, numpy_err in test_cases: a = torch.randn(input_size, device=device) - msg = f'input_size {input_size}, sections_or_indices {sections_or_indices}, dim {dim}' + msg = f"input_size {input_size}, sections_or_indices {sections_or_indices}, dim {dim}" with self.assertRaisesRegex(err, err_msg, msg=msg): torch.tensor_split(a, sections_or_indices, dim) with self.assertRaisesRegex(err, err_msg, msg=msg): @@ -1782,13 +1970,17 @@ class TestOldViewOps(TestCase): np.array_split(a.cpu().numpy(), sections_or_indices, dim) # addtional tests for tensor_split with tensor_indices_or_sections - with self.assertRaisesRegex(RuntimeError, - r'tensor_split expected tensor_indices_or_sections to have dtype of long, but got Float'): + with self.assertRaisesRegex( + RuntimeError, + r"tensor_split expected tensor_indices_or_sections to have dtype of long, but got Float", + ): torch.tensor_split(a, torch.tensor(1.1), dim) - with self.assertRaisesRegex(RuntimeError, - r'tensor_split expected tensor_indices_or_sections to be a' - + ' zero-dimensional or one-dimensional tensor, but got a tensor with 2 dims'): + with self.assertRaisesRegex( + RuntimeError, + r"tensor_split expected tensor_indices_or_sections to be a" + + " zero-dimensional or one-dimensional tensor, but got a tensor with 2 dims", + ): torch.tensor_split(torch.rand(S, device=device), torch.tensor(((1,),)), 0) def test_resize_all_dtypes_and_devices(self, device): @@ -1808,11 +2000,13 @@ class TestOldViewOps(TestCase): @onlyNativeDeviceTypes def test_resize_overflow(self, device): x = torch.empty((), dtype=torch.float64) - with self.assertRaisesRegex(RuntimeError, 'Storage size calculation overflowed'): + with self.assertRaisesRegex( + RuntimeError, "Storage size calculation overflowed" + ): x.resize_([2, 4, 2**29, 2**29]) - with self.assertRaisesRegex(RuntimeError, 'overflow'): + with self.assertRaisesRegex(RuntimeError, "overflow"): x.resize_([8, 8, 2**29, 2**29]) - with self.assertRaisesRegex(RuntimeError, 'Stride calculation overflowed'): + with self.assertRaisesRegex(RuntimeError, "Stride calculation overflowed"): x.resize_([0, 4, 2305843009213693952]) def test_view_all_dtypes_and_devices(self, device): @@ -1823,12 +2017,26 @@ class TestOldViewOps(TestCase): @skipIfTorchDynamo("conj bit not implemented in TensorVariable yet") @onlyCPU def test_conj_neg_view_numpy_error(self, device): - self.assertRaisesRegex(RuntimeError, "has conjugate bit set", lambda: torch.tensor([1 + 2j]).conj().numpy()) - self.assertRaisesRegex(RuntimeError, "has negative bit set", lambda: torch.tensor([1 + 2j]).conj().imag.numpy()) - self.assertRaisesRegex(RuntimeError, "not supported for conjugate view tensors", - lambda: torch.tensor([1 + 2j]).conj().view(torch.float64)) - self.assertRaisesRegex(RuntimeError, "not supported for tensors with negative bit set", - lambda: torch.tensor([1 + 2j]).conj().imag.view(torch.int32)) + self.assertRaisesRegex( + RuntimeError, + "has conjugate bit set", + lambda: torch.tensor([1 + 2j]).conj().numpy(), + ) + self.assertRaisesRegex( + RuntimeError, + "has negative bit set", + lambda: torch.tensor([1 + 2j]).conj().imag.numpy(), + ) + self.assertRaisesRegex( + RuntimeError, + "not supported for conjugate view tensors", + lambda: torch.tensor([1 + 2j]).conj().view(torch.float64), + ) + self.assertRaisesRegex( + RuntimeError, + "not supported for tensors with negative bit set", + lambda: torch.tensor([1 + 2j]).conj().imag.view(torch.int32), + ) @onlyCPU def test_crow_col_indices(self, device): @@ -1842,8 +2050,9 @@ class TestOldViewOps(TestCase): t.crow_indices() t.col_indices() + instantiate_device_type_tests(TestViewOps, globals(), include_lazy=True) instantiate_device_type_tests(TestOldViewOps, globals()) -if __name__ == '__main__': +if __name__ == "__main__": run_tests()