(self, device)
| 69 | |
| 70 | @skipIfRocm |
| 71 | def test_load_tensor(self, device): |
| 72 | with tempfile.TemporaryDirectory() as loc: |
| 73 | writer = ContentStoreWriter(loc) |
| 74 | x = torch.randn(4, device=device) |
| 75 | |
| 76 | def same_meta_as_x(t): |
| 77 | self.assertEqual(t.size(), x.size()) |
| 78 | self.assertEqual(t.stride(), x.stride()) |
| 79 | self.assertEqual(t.dtype, x.dtype) |
| 80 | self.assertEqual(t.device, x.device) |
| 81 | |
| 82 | writer.write_tensor("x", x) |
| 83 | |
| 84 | with load_tensor_reader(loc): |
| 85 | x2 = torch.ops.debugprims.load_tensor.default( |
| 86 | "x", (4,), (1,), dtype=torch.float32, device=device |
| 87 | ) |
| 88 | self.assertEqual(x, x2) |
| 89 | x3 = torch.ops.debugprims.load_tensor.default( |
| 90 | "x", (4,), (1,), dtype=torch.float32, device=device |
| 91 | ) |
| 92 | self.assertEqual(x, x3) |
| 93 | # Must not alias! |
| 94 | self.assertNotEqual( |
| 95 | StorageWeakRef(x.untyped_storage()), |
| 96 | StorageWeakRef(x2.untyped_storage()), |
| 97 | ) |
| 98 | self.assertNotEqual( |
| 99 | StorageWeakRef(x2.untyped_storage()), |
| 100 | StorageWeakRef(x3.untyped_storage()), |
| 101 | ) |
| 102 | |
| 103 | # Check fake tensor mode works too |
| 104 | with FakeTensorMode(): |
| 105 | x4 = torch.ops.debugprims.load_tensor.default( |
| 106 | "x", (4,), (1,), dtype=torch.float32, device=device |
| 107 | ) |
| 108 | self.assertIsInstance(x4, FakeTensor) |
| 109 | same_meta_as_x(x4) |
| 110 | |
| 111 | # Check fp64 works |
| 112 | x5 = torch.ops.debugprims.load_tensor.default( |
| 113 | "x", (4,), (1,), dtype=torch.float64, device=device |
| 114 | ) |
| 115 | self.assertEqual(x5.float(), x) |
| 116 | self.assertEqual(x5.dtype, torch.float64) |
| 117 | |
| 118 | x6 = torch.ops.debugprims.load_tensor.default( |
| 119 | "x", (4,), (1,), dtype=torch.float32, device=device |
| 120 | ) |
| 121 | same_meta_as_x(x6) |
| 122 | |
| 123 | |
| 124 | instantiate_device_type_tests(TestContentStore, globals()) |
nothing calls this directly
no test coverage detected