import torch
defmake_weight_cp(t, wa, wb):
temp = torch.einsum('i j k l, j r -> i r k l', t, wb)return torch.einsum('i j k l, i r -> r j k l', temp, wa)defrebuild_conventional(up, down, shape, dyn_dim=None):
up = up.reshape(up.size(0), -1)
down = down.reshape(down.size(0), -1)if dyn_dim is not None:
up = up[:, :dyn_dim]
down = down[:dyn_dim, :]return(up @ down).reshape(shape)