| import torch | |
| from peft.tuners._buffer_dict import BufferDict | |
| class TestBufferDict: | |
| def test_init_from_dict_works(self): | |
| bd = BufferDict( | |
| { | |
| "default": torch.randn(10, 2), | |
| } | |
| ) | |
| def test_update_from_other_bufferdict(self): | |
| default_tensor = torch.randn(10, 2) | |
| non_default_tensor = torch.randn(10, 2) | |
| bd1 = BufferDict({"default": default_tensor}) | |
| bd2 = BufferDict({"non_default": non_default_tensor}) | |
| bd1.update(bd2) | |
| assert set(bd1.keys()) == {"default", "non_default"} | |
| assert torch.allclose(bd1["default"], default_tensor) | |
| assert torch.allclose(bd1["non_default"], non_default_tensor) | |
| def test_update_from_dict(self): | |
| default_tensor = torch.randn(10, 2) | |
| non_default_tensor = torch.randn(10, 2) | |
| bd1 = BufferDict({"default": default_tensor}) | |
| d1 = {"non_default": non_default_tensor} | |
| bd1.update(d1) | |
| assert set(bd1.keys()) == {"default", "non_default"} | |
| assert torch.allclose(bd1["default"], default_tensor) | |
| assert torch.allclose(bd1["non_default"], non_default_tensor) | |
| def test_update_from_dict_items(self): | |
| default_tensor = torch.randn(10, 2) | |
| non_default_tensor = torch.randn(10, 2) | |
| bd1 = BufferDict({"default": default_tensor}) | |
| d1 = {"non_default": non_default_tensor} | |
| bd1.update(d1.items()) | |
| assert set(bd1.keys()) == {"default", "non_default"} | |
| assert torch.allclose(bd1["default"], default_tensor) | |
| assert torch.allclose(bd1["non_default"], non_default_tensor) | |