feat(autograd): add primitive for np.unwrap

Author: yaugenst-flex

CHANGELOG.md

@@ -7,6 +7,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+- Add support for `np.unwrap` in `tidy3d.plugins.autograd`.
+
 ### Fixed
 - Arrow lengths are now scaled consistently in the X and Y directions,
   and their lengths no longer exceed the height of the plot window.

tests/test_plugins/autograd/primitives/test_misc.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import autograd.numpy as np
 import pytest
 from autograd.test_util import check_grads
 
@@ -22,3 +23,19 @@
 def test_gaussian_filter_grad(rng, size, ndim, sigma, mode):
     x = rng.random((size,) * ndim)
     check_grads(lambda x: gaussian_filter(x, sigma=sigma, mode=mode), modes=["rev"], order=2)(x)
+
+
+@pytest.mark.parametrize("shape, axis", [((100,), -1), ((10, 12), 0), ((10, 12), 1)])
+@pytest.mark.parametrize("period", [np.pi, 2 * np.pi])
+@pytest.mark.parametrize("discont", [None, 0.6])
+def test_unwrap_grad(rng, shape, axis, period, discont):
+    """Test the gradient of the unwrap function with various arguments."""
+    if discont is not None:
+        # discont must be > period / 2 to have an effect
+        discont = discont * period
+
+    x = rng.uniform(-4 * period, 4 * period, shape)
+
+    check_grads(
+        lambda x: np.unwrap(x, discont=discont, axis=axis, period=period), modes=["fwd", "rev"]
+    )(x)

tidy3d/plugins/autograd/primitives/misc.py

@@ -1,10 +1,15 @@
 from __future__ import annotations
 
+import autograd.numpy as np
 import scipy.ndimage
-from autograd.extend import defvjp, primitive
+from autograd.extend import defjvp, defvjp, primitive
 
 gaussian_filter = primitive(scipy.ndimage.gaussian_filter)
 defvjp(
     gaussian_filter,
     lambda ans, x, *args, **kwargs: lambda g: gaussian_filter(g, *args, **kwargs),
 )
+
+np.unwrap = primitive(np.unwrap)
+defjvp(np.unwrap, lambda g, ans, x, *args, **kwargs: g)
+defvjp(np.unwrap, lambda ans, x, *args, **kwargs: lambda g: g)