Merge pull request #76 from facebookresearch/stop_grad

Add stop gradient to tangent vector field calcuation

Author: mfschubert

.bumpversion.toml

@@ -1,5 +1,5 @@
 [tool.bumpversion]
-current_version = "v0.5.2"
+current_version = "v0.5.3"
 commit = true
 commit_args = "--no-verify"
 tag = true

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 
 name = "fmmax"
-version = "v0.5.2"
+version = "v0.5.3"
 description = "Fourier modal method with Jax"
 readme = "README.md"
 requires-python = ">=3.7"

src/fmmax/__init__.py

@@ -1,6 +1,6 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 
-__version__ = "v0.5.2"
+__version__ = "v0.5.3"
 
 from . import (
     basis,

src/fmmax/vector.py

@@ -144,6 +144,7 @@ def compute_tangent_field(
     Returns:
         The normal field, `(tx, ty)`.
     """
+    arr = jax.lax.stop_gradient(arr)
     batch_shape = arr.shape[:-2]
     arr = utils.atleast_nd(arr, n=3)
     arr = arr.reshape((-1,) + arr.shape[-2:])