Update _lax_linalg.py

chaobrain · Jan 14, 2025 · f519bc6 · f519bc6
1 parent 7074057
commit f519bc6
Showing 1 changed file with 25 additions and 10 deletions.
diff --git a/brainunit/lax/_lax_linalg.py b/brainunit/lax/_lax_linalg.py
@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import sys
 from typing import Union, Callable, Any
 
 import jax
@@ -362,18 +363,32 @@ def svd(
     containing the left singular vectors, the singular values and the adjoint of
     the right singular vectors.
     """
-    if isinstance(x, Quantity):
-        if compute_uv:
-            u, s, vh = lax.linalg.svd(x.mantissa, full_matrices=full_matrices, compute_uv=compute_uv,
-                                      subset_by_index=subset_by_index, algorithm=algorithm)
-            return u, maybe_decimal(Quantity(s, unit=x.unit)), vh
+    if sys.version_info >= (3, 10):
+        if isinstance(x, Quantity):
+            if compute_uv:
+                u, s, vh = lax.linalg.svd(x.mantissa, full_matrices=full_matrices, compute_uv=compute_uv,
+                                          subset_by_index=subset_by_index, algorithm=algorithm)
+                return u, maybe_decimal(Quantity(s, unit=x.unit)), vh
+            else:
+                s = lax.linalg.svd(x.mantissa, full_matrices=full_matrices, compute_uv=compute_uv,
+                                   subset_by_index=subset_by_index, algorithm=algorithm)
+                return maybe_decimal(Quantity(s, unit=x.unit))
         else:
-            s = lax.linalg.svd(x.mantissa, full_matrices=full_matrices, compute_uv=compute_uv,
-                               subset_by_index=subset_by_index, algorithm=algorithm)
-            return maybe_decimal(Quantity(s, unit=x.unit))
+            return lax.linalg.svd(x, full_matrices=full_matrices, compute_uv=compute_uv,
+                                  subset_by_index=subset_by_index, algorithm=algorithm)
     else:
-        return lax.linalg.svd(x, full_matrices=full_matrices, compute_uv=compute_uv,
-                              subset_by_index=subset_by_index, algorithm=algorithm)
+        if isinstance(x, Quantity):
+            if compute_uv:
+                u, s, vh = lax.linalg.svd(x.mantissa, full_matrices=full_matrices, compute_uv=compute_uv,
+                                          subset_by_index=subset_by_index)
+                return u, maybe_decimal(Quantity(s, unit=x.unit)), vh
+            else:
+                s = lax.linalg.svd(x.mantissa, full_matrices=full_matrices, compute_uv=compute_uv,
+                                   subset_by_index=subset_by_index)
+                return maybe_decimal(Quantity(s, unit=x.unit))
+        else:
+            return lax.linalg.svd(x, full_matrices=full_matrices, compute_uv=compute_uv,
+                                  subset_by_index=subset_by_index)
 
 
 @set_module_as('brainunit.lax')