clean up jax svd version

PlasmaControl · Nov 21, 2024 · b116a8a · b116a8a
1 parent ca3d967
commit b116a8a
Showing 1 changed file with 1 addition and 3 deletions.
diff --git a/desc/utils.py b/desc/utils.py
@@ -424,9 +424,7 @@ def svd_inv_null(A):
     num = jnp.sum(large, dtype=int)
     uk = u[:, :K]
     vhk = vh[:K, :]
-    s = jnp.where(large, 1 / s, s)
-    s.shape
-    s = s.at[(~large,)].set(0)
+    s = jnp.where(large, 1 / s, 0)
     Ainv = vhk.T @ jnp.diag(s) @ uk.T
     Z = vh[num:, :].T.conj()
     return Ainv, Z