1import json, math, random, time
  2import numpy as np
  3import torch
  4import torch.nn as nn
  5import torch.nn.functional as F
  6
  7SEED = 605
  8np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
  9
 10
 11def tt_svd_matrix(W, modes, rel_tol=0.0, max_rank=None):
 12    """TT-SVD of W with C-order input/output mode ordering."""
 13    m, n = modes
 14    d = len(m)
 15    T = W.reshape(*(list(m) + list(n)))
 16    # Convert to interleaved physical ordering (m1,n1,m2,n2,...).
 17    perm = [k for pair in zip(range(d), range(d, 2*d)) for k in pair]
 18    T = T.transpose(perm)
 19    norm = np.linalg.norm(W)
 20    eps2 = (rel_tol * norm) ** 2
 21    # Equal per-bond budget gives the standard global error bound.
 22    budget = eps2 / max(1, d - 1)
 23    cores, ranks, tails = [], [1], []
 24    left = T
 25    rprev = 1
 26    for k in range(d - 1):
 27        left = left.reshape(rprev * m[k] * n[k], -1)
 28        u, s, vh = np.linalg.svd(left, full_matrices=False)
 29        tail = np.cumsum(s[::-1] ** 2)[::-1]
 30        rank = len(s)
 31        if rel_tol > 0:
 32            valid = np.where(np.r_[tail[1:], 0.0] <= budget + 1e-14)[0]
 33            if len(valid): rank = int(valid[0] + 1)
 34        if max_rank is not None: rank = min(rank, max_rank)
 35        rank = max(1, rank)
 36        discarded = float(np.sum(s[rank:] ** 2))
 37        tails.append(discarded)
 38        cores.append(u[:, :rank].reshape(rprev, m[k], n[k], rank))
 39        left = (s[:rank, None] * vh[:rank])
 40        rprev = rank; ranks.append(rank)
 41    cores.append(left.reshape(rprev, m[-1], n[-1], 1))
 42    ranks.append(1)
 43    return cores, ranks, tails, norm
 44
 45
 46def tt_reconstruct(cores):
 47    x = cores[0]
 48    for c in cores[1:]:
 49        x = np.tensordot(x, c, axes=([-1], [0]))
 50    # x: m1,n1,m2,n2,..., with singleton bond ends removed
 51    d = len(cores)
 52    x = np.squeeze(x, axis=(0, -1))
 53    shape = x.shape
 54    out_modes = [shape[2*k] for k in range(d)]
 55    in_modes = [shape[2*k+1] for k in range(d)]
 56    perm = list(range(0, 2*d, 2)) + list(range(1, 2*d, 2))
 57    return x.transpose(perm).reshape(int(np.prod(out_modes)), int(np.prod(in_modes)))
 58
 59
 60def tt_param_count(modes, ranks):
 61    m, n = modes
 62    return sum(ranks[k] * m[k] * n[k] * ranks[k+1] for k in range(len(m)))
 63
 64
 65class TTLinear(nn.Module):
 66    def __init__(self, W, modes, rel_tol=0.0, max_rank=None, bias=True):
 67        super().__init__(); self.modes = modes; self.max_rank = max_rank
 68        cores, ranks, _, _ = tt_svd_matrix(W, modes, rel_tol, max_rank)
 69        self.cores = nn.ParameterList([nn.Parameter(torch.tensor(c, dtype=torch.float32)) for c in cores])
 70        self.ranks = ranks
 71        self.bias = nn.Parameter(torch.zeros(W.shape[0])) if bias else None
 72
 73    def forward(self, x):
 74        # Direct tensor-network contraction: batch, all input modes, TT cores.
 75        d = len(self.modes[0])
 76        letters = list("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ")
 77        out = letters[1:1+d]; inn = letters[1+d:1+2*d]
 78        bond = letters[1+2*d:1+3*d+1]  # includes both singleton boundary bonds
 79        xsub = letters[0] + ''.join(inn)
 80        terms = [xsub]
 81        for k, c in enumerate(self.cores):
 82            terms.append(bond[k] + out[k] + inn[k] + bond[k+1])
 83        equation = ','.join(terms) + '->' + letters[0] + ''.join(out)
 84        operands = [x.reshape(x.shape[0], *self.modes[1])] + list(self.cores)
 85        y = torch.einsum(equation, *operands).reshape(x.shape[0], -1)
 86        return y + self.bias if self.bias is not None else y
 87
 88    @torch.no_grad()
 89    def round(self, rel_tol):
 90        W = tt_reconstruct([c.detach().cpu().numpy() for c in self.cores])
 91        cores, ranks, _, _ = tt_svd_matrix(W, self.modes, rel_tol, self.max_rank)
 92        self.cores = nn.ParameterList([nn.Parameter(torch.tensor(c, device=self.cores[0].device, dtype=self.cores[0].dtype)) for c in cores])
 93        self.ranks = ranks
 94
 95
 96def verify():
 97    modes = ([4,4,4], [4,4,4]); W = np.random.randn(64,64)
 98    cores, ranks, tails, norm = tt_svd_matrix(W, modes, rel_tol=0.08, max_rank=8)
 99    Wa = tt_reconstruct(cores); err = np.linalg.norm(W-Wa)
100    # Uncapped TT-SVD tests the stated tolerance theorem directly.
101    exact_cores, exact_ranks, exact_tails, _ = tt_svd_matrix(W, modes, rel_tol=0.08, max_rank=None)
102    exact_err = np.linalg.norm(W - tt_reconstruct(exact_cores))
103    bound = 0.08 * norm
104    # contraction equality check on a random vector
105    layer = TTLinear(W, modes, rel_tol=0.08, max_rank=8, bias=False)
106    x = torch.randn(3,64)
107    direct = layer(x).detach().numpy(); expected = x.numpy() @ Wa.T
108    return {'relative_fro_error': float(err/norm), 'requested_bound': float(bound/norm),
109            'bound_holds_with_cap': bool(err <= bound + 1e-6),
110            'uncapped_relative_fro_error': float(exact_err/norm),
111            'uncapped_bound_holds': bool(exact_err <= bound + 1e-6),
112            'contraction_max_abs': float(np.max(np.abs(direct-expected))),
113            'ranks': ranks, 'uncapped_ranks': exact_ranks, 'dense_params': 4096,
114            'tt_params': tt_param_count(modes, ranks), 'uncapped_tt_params': tt_param_count(modes, exact_ranks),
115            'tail_squared_sum': float(sum(tails))}
116
117
118def train_compare():
119    device = 'cuda' if torch.cuda.is_available() else 'cpu'
120    try:
121        torch.manual_seed(SEED)
122        modes = ([4,4,4], [4,4,4]); dim=64
123        # Teacher is genuinely TT-low-rank, making this a representation test.
124        teacher = TTLinear(np.random.randn(dim,dim)*0.25, modes, rel_tol=0, max_rank=2, bias=True).to(device)
125        with torch.no_grad(): teacher.bias.normal_(0, .1)
126        X = torch.randn(1024, dim, device=device); Y = torch.tanh(teacher(X)).detach()
127        results = {}
128        for name, model in [('dense', nn.Sequential(nn.Linear(dim,dim), nn.Tanh()).to(device)),
129                            ('tt', None)]:
130            if model is None:
131                model = nn.Sequential(TTLinear(np.random.randn(dim,dim)*.05, modes, rel_tol=0, max_rank=4).to(device), nn.Tanh())
132            opt = torch.optim.Adam(model.parameters(), lr=3e-3)
133            t0=time.perf_counter(); losses=[]
134            for step in range(250):
135                pred=model(X); loss=F.mse_loss(pred,Y); opt.zero_grad(); loss.backward(); opt.step()
136                if name=='tt' and (step+1)%50==0: model[0].round(0.03)
137                losses.append(float(loss.detach().cpu()))
138            results[name]={'final_mse':losses[-1], 'mse_step_50':losses[49], 'seconds':time.perf_counter()-t0,
139                           'params':sum(p.numel() for p in model.parameters()),
140                           'ranks': getattr(model[0], 'ranks', None) if name=='tt' else None}
141        return {'device':device, 'results':results}
142    except Exception as e:
143        if device == 'cuda':
144            torch.cuda.empty_cache(); return train_compare_cpu()
145        raise
146
147def train_compare_cpu():
148    torch.cuda.is_available=lambda: False
149    return train_compare()
150
151if __name__ == '__main__':
152    print(json.dumps({'verification':verify(), 'training':train_compare()}, indent=2))