File size: 1,500 Bytes
3236af9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 | import sys
import unittest
from pathlib import Path
import torch
sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))
from iconoclast.direction import (
compute_benign_subspace_basis,
project_directions_out_of_subspace,
)
class DirectionTests(unittest.TestCase):
def test_benign_subspace_projection_removes_principal_good_direction(self):
good_residuals = torch.tensor(
[
[[3.0, 0.0, 0.0]],
[[1.0, 0.0, 0.0]],
[[-1.0, 0.0, 0.0]],
[[-3.0, 0.0, 0.0]],
]
)
basis = compute_benign_subspace_basis(good_residuals, rank=1)
directions = torch.tensor([[1.0, 1.0, 0.0]])
projected = project_directions_out_of_subspace(directions, basis)
self.assertAlmostEqual(projected[0, 0].item(), 0.0, places=5)
self.assertAlmostEqual(projected[0, 1].item(), 1.0, places=5)
self.assertAlmostEqual(projected[0, 2].item(), 0.0, places=5)
def test_zero_rank_benign_subspace_is_disabled(self):
good_residuals = torch.tensor(
[
[[1.0, 0.0]],
[[-1.0, 0.0]],
]
)
basis = compute_benign_subspace_basis(good_residuals, rank=0)
directions = torch.tensor([[0.6, 0.8]])
projected = project_directions_out_of_subspace(directions, basis)
self.assertTrue(torch.equal(projected, directions))
if __name__ == "__main__":
unittest.main()
|