fix: xarray 2023.08.0 compatibility

Author: yaugenst-flex

tests/test_components/test_autograd.py

@@ -1630,7 +1630,7 @@ def test_custom_methods_grads(self, attr):
         """Test grads of TidyArrayBox methods implemented in autograd/boxes.py"""
 
         def objective(x, attr):
-            da = DataArray(x, dims=map(str, range(x.ndim)))
+            da = DataArray(x)
             attr_value = getattr(da, attr)
             val = attr_value() if callable(attr_value) else attr_value
             return val.item()
@@ -1643,11 +1643,11 @@ def test_multiply_at_grads(self, rng):
 
         def objective(a, b):
             coords = {str(i): np.arange(a.shape[i]) for i in range(a.ndim)}
-            da = DataArray(a, coords=coords, dims=map(str, range(a.ndim)))
+            da = DataArray(a, coords=coords)
             da_mult = da.multiply_at(b, "0", [0, 1]) ** 2
             return np.sum(da_mult).item()
 
         a = rng.uniform(-1, 1, (3, 3))
         b = 1.0
-        check_grads(lambda x: objective(x, b), modes=["fwd", "rev"], order=1)(a)
-        check_grads(lambda x: objective(a, x), modes=["fwd", "rev"], order=1)(b)
+        check_grads(lambda x: objective(x, b), modes=["fwd", "rev"], order=2)(a)
+        check_grads(lambda x: objective(a, x), modes=["fwd", "rev"], order=2)(b)

tidy3d/components/autograd/boxes.py

@@ -5,12 +5,21 @@
 from typing import Any, Callable, Dict, List, Tuple
 
 import autograd.numpy as anp
+from autograd.extend import defjvp
 from autograd.numpy.numpy_boxes import ArrayBox
+from autograd.numpy.numpy_wrapper import _astype
 
 TidyArrayBox = ArrayBox  # NOT a subclass
 
 _autograd_module_cache = {}  # cache for imported autograd modules
 
+defjvp(
+    _astype,
+    lambda g, ans, A, dtype, order="K", casting="unsafe", subok=True, copy=True: _astype(g, dtype),
+)
+
+anp.astype = _astype
+
 
 @classmethod
 def from_arraybox(cls, box: ArrayBox) -> TidyArrayBox: