fix: xarray 2024.10 compatibility

Author: yaugenst-flex

CHANGELOG.md

@@ -17,6 +17,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Fixed
 - Regression in local field projection leading to incorrect results for `far_field_approx=True`.
 - Bug when differentiating with respect to `Cylinder.center`.
+- `xarray` 2024.10.0 compatibility for autograd.
 
 
 ## [2.7.6] - 2024-10-30

poetry.lock

@@ -57,7 +57,7 @@ ipython = { version = "*", optional = true }
 memory_profiler = { version = "*", optional = true }
 pre-commit = { version = "*", optional = true }
 pylint = { version = "*", optional = true }
-pytest = { version = "*", optional = true }
+pytest = { version = ">=8.1", optional = true }
 pytest-timeout = { version = "*", optional = true }
 tox = { version = "*", optional = true }
 

pyproject.toml

@@ -1,6 +1,3 @@
-import autograd.numpy as anp
-from autograd.extend import VJPNode, register_notrace
-
 from .boxes import TidyArrayBox
 from .functions import interpn
 from .types import (
@@ -14,8 +11,6 @@
 )
 from .utils import get_static, is_tidy_box, split_list
 
-register_notrace(VJPNode, anp.full_like)
-
 __all__ = [
     "TidyArrayBox",
     "TracedFloat",

tidy3d/components/autograd/__init__.py

@@ -1,24 +1,27 @@
-# Adds some functionality to the autograd arraybox
-# NOTE: this is not a subclass of ArrayBox since that would break autograd's internal checks
+# Adds some functionality to the autograd arraybox and related autograd patches
+# NOTE: we do not subclass ArrayBox since that would break autograd's internal checks
 
 import importlib
 from typing import Any, Callable, Dict, List, Tuple
 
 import autograd.numpy as anp
-from autograd.extend import defjvp
+from autograd.extend import VJPNode, defjvp, register_notrace
 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
 
+register_notrace(VJPNode, anp.full_like)
+
 defjvp(
     _astype,
     lambda g, ans, A, dtype, order="K", casting="unsafe", subok=True, copy=True: _astype(g, dtype),
 )
 
 anp.astype = _astype
+anp.permute_dims = anp.transpose
 
 
 @classmethod