clean up vectors_yx

Author: Jammy2211

autoarray/structures/vectors/uniform.py

@@ -1,7 +1,8 @@
 import logging
 
-# import numpy as np
-from autofit.jax_wrapper import numpy as np, use_jax
+import numpy as np
+import jax.numpy as jnp
+# from autofit.jax_wrapper import numpy as np, use_jax
 from typing import List, Optional, Tuple, Union
 
 from autoarray.structures.arrays.uniform_2d import Array2D
@@ -396,11 +397,7 @@ def magnitudes(self) -> Array2D:
         """
         Returns the magnitude of every vector which are computed as sqrt(y**2 + x**2).
         """
-        if use_jax:
-            s = self.array
-        else:
-            s = self
-        return Array2D(values=np.sqrt(s[:, 0] ** 2.0 + s[:, 1] ** 2.0), mask=self.mask)
+        return Array2D(values=jnp.sqrt(self.array[:, 0] ** 2.0 + self.array[:, 1] ** 2.0), mask=self.mask)
 
     @property
     def y(self) -> Array2D: