When using JAX don't import Numpy

Switches out `numpy` for `jax.numpy`.
Jammy2211 · Jun 28, 2024 · 044febb · 044febb
1 parent 7f0bcec
commit 044febb
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/autogalaxy/profiles/geometry_profiles.py b/autogalaxy/profiles/geometry_profiles.py
@@ -1,6 +1,11 @@
+import os
+
 from typing import Optional, Tuple, Type
 
-import numpy as np
+if os.environ.get("USE_JAX", "0") == "1":
+    import jax.numpy as np
+else:
+    import numpy as np
 
 import autoarray as aa