| 272 | |
| 273 | |
| 274 | def test_numpy_save_aligned(): |
| 275 | assert_equal = partial(torch.testing.assert_close, rtol=0, atol=0) |
| 276 | a = torch.randn(1024, dtype=torch.float32) # 4096 bytes |
| 277 | with tempfile.TemporaryDirectory() as test_dir: |
| 278 | aligned_path = os.path.join(test_dir, "aligned.npy") |
| 279 | gb.numpy_save_aligned(aligned_path, a.numpy()) |
| 280 | |
| 281 | nonaligned_path = os.path.join(test_dir, "nonaligned.npy") |
| 282 | np.save(nonaligned_path, a.numpy()) |
| 283 | |
| 284 | assert_equal(np.load(aligned_path), np.load(nonaligned_path)) |
| 285 | # The size of the file should be 4K (aligned header) + 4K (tensor). |
| 286 | assert os.path.getsize(aligned_path) == 4096 * 2 |