Merge pull request #76 from facebookresearch/stop_grad

Add stop gradient to tangent vector field calcuation
facebookresearch · Jan 11, 2024 · 8c7491a · 8c7491a
2 parents d461c4b + f75a975
commit 8c7491a
Show file tree

Hide file tree

Showing 4 changed files with 4 additions and 3 deletions.
diff --git a/.bumpversion.toml b/.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

diff --git a/pyproject.toml b/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"

diff --git a/src/fmmax/__init__.py b/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,

diff --git a/src/fmmax/vector.py b/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:])