Fix/sinkdiv hessian (#321)

* Use symmetric in `sinkhorn_divergence`

* Re-run notebook

* [ci skip] Regenerate notebook

src/ott/solvers/linear/implicit_differentiation.py

@@ -12,8 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Functions entering the implicit differentiation of Sinkhorn."""
-
-from typing import TYPE_CHECKING, Callable, Optional, Tuple
+import dataclasses
+from typing import TYPE_CHECKING, Any, Callable, Optional, Tuple
 
 import jax
 import jax.numpy as jnp
@@ -286,3 +286,6 @@ def gradient(
     # Carries pullback onto original inputs, here geom, a and b.
     _, pull_prob = jax.vjp(foc_prob, prob)
     return pull_prob(vjp_gr)
+
+  def replace(self, **kwargs: Any) -> "ImplicitDiff":  # noqa: D102
+    return dataclasses.replace(self, **kwargs)

src/ott/tools/sinkhorn_divergence.py

@@ -159,8 +159,10 @@ def _sinkhorn_divergence(
         parallel_dual_updates=True,
         momentum=acceleration.Momentum(start=0, value=0.5),
         anderson=None,
-        # TODO(michalk8): implicit_diff
     )
+    implicit_diff = kwargs.get("implicit_diff", None)
+    if implicit_diff is not None:
+      kwargs_symmetric["implicit_diff"] = implicit_diff.replace(symmetric=True)
 
   out_xy = sinkhorn.solve(geometry_xy, a, b, **kwargs)
   out_xx = sinkhorn.solve(geometry_xx, a, a, **kwargs_symmetric)