code: def test_custom_gradient_squared_dist(self):
      def squared_dist_grad_a(point_a, point_b, metric):
          return -2 * metric.log(point_b, point_a)
  
      def squared_dist_grad_b(point_a, point_b, metric):
          return -2 * metric.log(point_a, point_b)
  
      @gs.autodiff.custom_gradient(
          squared_dist_grad_a, squared_dist_grad_b
      )
      def squared_dist(point_a, point_b, metric):
          dist = metric.squared_dist(point_a, point_b)
          return dist
  
      space = SpecialEuclidean(n=2)
      const_metric = space.metric
      const_point_b = space.random_point()
  
      def func(x):
          return squared_dist(x, const_point_b, metric=const_metric)
  
      arg_point_a = space.random_point()
      expected_value = func(arg_point_a)
      expected_grad = -2 * const_metric.log(const_point_b, arg_point_a)
      result_value, result_grad = gs.autodiff.value_and_grad(func)(
          arg_point_a
      )
      self.assertAllClose(result_value, expected_value)
      self.assertAllClose(result_grad, expected_grad)
  
imports: import warnings
  import numpy
  import pytest
  import geomstats.backend
  from geomstats.geometry.grassmannian import Grassmannian
  from geomstats.geometry.special_euclidean import SpecialEuclidean
  from geomstats.test.test_case import TestCase, autograd_and_torch_only, autograd_only, np_only
label_desc: Use image node_id 17 for calling the TestAutodiff obj's underlying member method code with example usage: obj.test_custom_gradient_squared_dist() and returns: unknown, unknown, dist, squared_dist