FIX: switch back to jacfwd

ComPWA · Aug 11, 2023 · 922a85f · 922a85f
1 parent b5aed2d
commit 922a85f
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/.cspell.json b/.cspell.json
@@ -213,7 +213,7 @@
     "ipywidgets",
     "isinstance",
     "isort",
-    "jacrev",
+    "jacfwd",
     "jaxlib",
     "joinpath",
     "juliaup",

diff --git a/docs/report/draft.ipynb b/docs/report/draft.ipynb
@@ -492,7 +492,7 @@
    "outputs": [],
    "source": [
     "func_with_data_inserted = Partial(intensity_func.function, *data_columns.values())\n",
-    "gradient_func = jax.jacrev(\n",
+    "gradient_func = jax.jacfwd(\n",
     "    func_with_data_inserted,\n",
     "    argnums=range(len(parameter_values)),\n",
     ")\n",
@@ -578,7 +578,7 @@
     "import jax.numpy as jnp\n",
     "\n",
     "\n",
-    "# @jax.jit  # Do not JIT here, otherwise jax.jacrev crashes!\n",
+    "# @jax.jit  # Do not JIT here, otherwise jax.jacfwd crashes!\n",
     "def estimator(args):\n",
     "    data_intensities = func_with_data_inserted(*args)\n",
     "    phsp_intensities = func_with_phsp_inserted(*args)\n",
@@ -640,7 +640,7 @@
    },
    "outputs": [],
    "source": [
-    "estimator_gradient = jax.jacrev(estimator)"
+    "estimator_gradient = jax.jacfwd(estimator)"
    ]
   },
   {