From d03dcedfdba0b31a6be470e9fdb69f7fde6932d0 Mon Sep 17 00:00:00 2001 From: Gabriel Kenneth Kosmacher Date: Fri, 19 Sep 2025 08:46:19 -0600 Subject: [PATCH 1/2] MTTKRP speedup (1.5x - 6x range) with einsum Replace the Y, Ur loop ``` V = np.zeros((szn, R)) for r in range(R): V[:, [r]] = Y[:, :, r].T @ Ur[:, :, r] ``` with ``` V = np.einsum('ijk, ik -> jk', Y, Ur) ``` for numpy vectorization with einsum. Some profiling results: shape [20, 30, 40, 50] --------------------- mode-1 mttkrp rank-10 new mttkrp correct True old 0.00010779168870713975 new 6.477038065592448e-05 speedup 1.6642126789366052 ----------------------------------------------------------------------- rank-50 new mttkrp correct True old 0.0006463527679443359 new 0.000177409913804796 speedup 3.6432731073615052 ----------------------------------------------------------------------- rank-100 new mttkrp correct True old 0.0006452401479085287 new 0.00010201666090223525 speedup 6.324850688132952 ----------------------------------------------------------------------- rank-200 new mttkrp correct True old 0.0016609827677408855 new 0.0002986590067545573 speedup 5.561468866418307 ----------------------------------------------------------------------- mode-1 mttkrp rank-10 new mttkrp correct True old 0.0004661348130967882 new 0.00024816724989149306 speedup 1.8783091374893253 ----------------------------------------------------------------------- rank-50 new mttkrp correct True old 0.002485460705227322 new 0.001473824183146159 speedup 1.6864024445043586 ----------------------------------------------------------------------- rank-100 new mttkrp correct True old 0.006084998448689778 new 0.0025424957275390625 speedup 2.3933170792698175 ----------------------------------------------------------------------- rank-200 new mttkrp correct True old 0.04606061511569553 new 0.02842871348063151 speedup 1.6202145463626638 ----------------------------------------------------------------------- --- pyttb/tensor.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/pyttb/tensor.py b/pyttb/tensor.py index 38cca47a..2ce818ae 100644 --- a/pyttb/tensor.py +++ b/pyttb/tensor.py @@ -1132,15 +1132,11 @@ def mttkrp( return to_memory_order(Y.T @ Ul, self.order) else: Ul = ttb.khatrirao(*U[n + 1 :], reverse=True) - Ur = np.reshape( - ttb.khatrirao(*U[0:n], reverse=True), (szl, 1, R), order=self.order - ) + Ur = ttb.khatrirao(*U[0:n], reverse=True) Y = np.reshape(self.data, (-1, szr), order=self.order) Y = Y @ Ul Y = np.reshape(Y, (szl, szn, R), order=self.order) - V = np.zeros((szn, R), order=self.order) - for r in range(R): - V[:, [r]] = Y[:, :, r].T @ Ur[:, :, r] + V = np.einsum('ijk, ik -> jk', Y, Ur) return to_memory_order(V, self.order) def mttkrps(self, U: ttb.ktensor | Sequence[np.ndarray]) -> list[np.ndarray]: From a35ec0aec0239f8358f6deac91931c9c1645d352 Mon Sep 17 00:00:00 2001 From: Gabriel Kosmacher <73120774+kennykos@users.noreply.github.com> Date: Mon, 22 Sep 2025 07:59:58 -0600 Subject: [PATCH 2/2] Update einsum string Co-authored-by: Nick <24689722+ntjohnson1@users.noreply.github.com> --- pyttb/tensor.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyttb/tensor.py b/pyttb/tensor.py index 2ce818ae..dd4e8030 100644 --- a/pyttb/tensor.py +++ b/pyttb/tensor.py @@ -1136,7 +1136,7 @@ def mttkrp( Y = np.reshape(self.data, (-1, szr), order=self.order) Y = Y @ Ul Y = np.reshape(Y, (szl, szn, R), order=self.order) - V = np.einsum('ijk, ik -> jk', Y, Ur) + V = np.einsum("ijk, ik -> jk", Y, Ur) return to_memory_order(V, self.order) def mttkrps(self, U: ttb.ktensor | Sequence[np.ndarray]) -> list[np.ndarray]: