tocmo0nlord/axolotl
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

add a prelim test for expading the 4d mask

Browse Source
This commit is contained in:
Wing Lian
2024-01-26 00:41:24 -05:00
parent 74c72ca5eb
commit 1a538be9c2
1 changed files with 26 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

26
tests/test_expand_mask.py
View File

@@ -39,6 +39,32 @@ class TestExpandMask(unittest.TestCase):
# Check that the output matches the expected output
self.assertTrue(torch.allclose(_expand_mask(mask, dtype), expected_output))
def test_output_multipack(self):
mask = torch.tensor([[1, 1, 1, 0], [2, 2, 3, 3]])
dtype = torch.float32
expected_output = torch.tensor(
[
[
[
[0.0000e00, -3.4028e38, -3.4028e38, -3.4028e38],
[0.0000e00, 0.0000e00, -3.4028e38, -3.4028e38],
[0.0000e00, 0.0000e00, 0.0000e00, -3.4028e38],
[-3.4028e38, -3.4028e38, -3.4028e38, -3.4028e38],
]
],
[
[
[0.0000e00, -3.4028e38, -3.4028e38, -3.4028e38],
[0.0000e00, 0.0000e00, -3.4028e38, -3.4028e38],
[-3.4028e38, -3.4028e38, 0.0000e00, -3.4028e38],
[-3.4028e38, -3.4028e38, 0.0000e00, 0.0000e00],
]
],
]
)
# Check that the output matches the expected output
self.assertTrue(torch.allclose(_expand_mask(mask, dtype), expected_output))
if __name__ == "__main__":
unittest.main()