avoid full matrices in svd init

mancusolab · Nov 11, 2022 · f8995c8 · f8995c8
1 parent 7f5c1d7
commit f8995c8
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/susiepca/infer.py b/susiepca/infer.py
@@ -146,7 +146,7 @@ def init_params(
 
     if init == "pca":
         # run PCA and extract weights and latent
-        svd_result = jnp.linalg.svd(X - jnp.mean(X, axis=0))
+        svd_result = jnp.linalg.svd(X - jnp.mean(X, axis=0), full_matrices=False)
         init_mu_z = svd_result[0] @ jnp.diag(svd_result[1])[:, 0:z_dim]
     elif init == "random":
         # random initialization