Add support for svd - #70
Conversation
gausshj
commented
Oct 9, 2025
gausshj
commented
- Add support for svd
- Add test for svd
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
hzhangxyz
left a comment
hzhangxyz
left a comment
There was a problem hiding this comment.
我们先review一下业务逻辑，先改这数值svd部分吧。你修改后，我们再说代码风格的事情。
数值svd后，恢复grassmann tensor这部分看起来没毛病。
|
|
||
| tensor = tensor.reshape((left_dim, right_dim)) | ||
|
|
||
| U, S, Vh = torch.linalg.svd(tensor.tensor, full_matrices=full_matrices) |
| U, S, Vh = torch.linalg.svd(tensor.tensor, full_matrices=full_matrices) | ||
|
|
||
| k = min(tensor.tensor.shape[0], tensor.tensor.shape[-1]) | ||
| k_index = tensor.tensor.shape.index(k) |
There was a problem hiding this comment.
分别进行svd后，需要允许有cut dimension的操作，这个在tn中很常见。大概就是删掉最小几个singular value，只保留最大的若干个，这个个数使用参数传进来，默认不进行cut，这里两个分块需要分别cut。
8bb8f29 to
e8297cb
Compare
hzhangxyz
left a comment
hzhangxyz
left a comment
There was a problem hiding this comment.
从你的代码看起来你是理解svd应该怎么做了，不过我们要做的是带broadcast的fermi tensor，所以要有一些不同。
| self, | ||
| free_names_u: tuple[int, ...], | ||
| *, | ||
| full_matrices: bool = False, # When full_matrices=True, the gradient with respect to U and Vh will be ignored |
There was a problem hiding this comment.
full matrics可以直接是False，tensor下面，不会用到full matrics的svd。
|
|
||
| tensor = tensor.reshape((left_dim, right_dim)) | ||
|
|
||
| tensor.update_mask() |
| keep_even = torch.ones_like(S_even, dtype=torch.bool, device=S_even.device) | ||
| keep_odd = torch.ones_like(S_odd, dtype=torch.bool, device=S_odd.device) | ||
| else: | ||
| S_cat = torch.cat([S_even, S_odd]) |
There was a problem hiding this comment.
你的想法是对的，最普适的svd确实是应该这么做，不过我们需要broadcast的svd，所以如果直接设置一个cutoff会使得每个batch出现不同的S even/S odd选择方案。所以我们应该在S even里选择最大的若干个，S odd里选择最大的若干个。他们之间互不影响。这个函数的cutoff类型可以是None | int | tuple[int, int]
e8297cb to
4db55dd
Compare
| odd_tensor = tensor.tensor[even_left:, even_right:] | ||
|
|
||
| if even_tensor.numel() > 0: | ||
| U_even, S_even, Vh_even = torch.linalg.svd(even_tensor, full_matrices=False) |
There was a problem hiding this comment.
svd出来的singular本来就是按着顺序排序的, 你不需要再重新排序
| n_even, n_odd = S_even.shape[0], S_odd.shape[0] | ||
| total = n_even + n_odd | ||
|
|
||
| if cutoff is None: |
There was a problem hiding this comment.
并不需要这么多判断的
if cutoff is None:
cutoff = (n_even, n_odd)
if isinstance(cutoff, int):
cutoff = (cutoff, cutoff)
cutoff = (n_even if cutoff[0] > n_even, cutoff[0], ...)
然后
U_even = U_even[:, :cutoff[0]]
S_even = S_even[:cutoff[0]]
Vh_even = Vh_even[:cutoff[0], :]
就行了
| cursor_plan: int = 0 | ||
| cursor_self: int = 0 | ||
| while cursor_plan != len(new_shape) or cursor_self != self.tensor.dim(): | ||
| if cursor_self == self.tensor.dim() and cursor_plan != len(new_shape): |
There was a problem hiding this comment.
算了, 你还是把这个部分的整理下, 不要用eo这种变量名, 做个临时的fix合进来吧
9fd8258 to
b9ada2d
Compare
- Add support for svd, perform SVD separately on even/odd parity blocks - Add support for two type of cutoff: int and tuple[int, int] - Add support for grassmann tensor with empty parity block - Add parameterized test cases for svd
b9ada2d to
3ad4a5b
Compare
2 participants
