cast abs(JaxDataArray) to jnp.array

CHANGELOG.md

@@ -38,6 +38,7 @@ the difference that can be observed when slightly modifying the grid resolution.
 - Numerically stable sigmoid function in radius of curvature constraint.
 - Fixed 2d checking in `Geometry.intersections_2dbox()`.
 - Spatial monitor downsampling when the monitor is crossing a symmetry plane or Bloch boundary conditions.
+- Cast `JaxDataArray.__abs__` output to `jnp.array`, reducing conversions needed in objective functions.
 
 ## [2.4.0rc1] - 2023-7-27
 

tidy3d/plugins/adjoint/components/data/data_array.py

@@ -163,7 +163,7 @@ def conj(self) -> JaxDataArray:
 
     def __abs__(self) -> JaxDataArray:
         """Absolute value of self's values."""
-        new_values = jnp.abs(self.values)
+        new_values = jnp.abs(self.as_jnp_array)
         return self.updated_copy(values=new_values)
 
     def __pow__(self, power: int) -> JaxDataArray:
@@ -213,7 +213,7 @@ def sum(self, dim: str = None):
         """Sum (optionally along a single or multiple dimensions)."""
 
         if dim is None:
-            return jnp.sum(self.values)
+            return jnp.sum(self.as_jnp_array)
 
         # dim is supplied
         if isinstance(dim, str):