code: def test_kahan_custom_grad(self):
      def model_fn():
          root = joint_distribution.JointDistributionCoroutine.Root
          _ = yield root(
              sharded.Sharded(
                  independent.Independent(
                      normal.Normal(0, tf.ones([7])),
                      reinterpreted_batch_ndims=1,
                      experimental_use_kahan_sum=True,
                  ),
                  shard_axis_name=self.axis_name,
              )
          )
  
      model = joint_distribution.JointDistributionCoroutine(model_fn)
  
      samps = self.strategy_run(
          lambda seed: model.sample(seed=seed),
          (test_util.test_seed(sampler_type="stateless"),),
          in_axes=None,
      )
  
      # TODO(b/204228837): This used to work with jit_compile=True, but we had to
      # disable it due to cross-device resource access errors.
      @tf.function(autograph=False)
      def lp_grad(x):
          return gradient.value_and_gradient(model.log_prob, x)
  
      self.evaluate(
          self.per_replica_to_tensor(
              self.strategy_run(lp_grad, (samps,))
          )
      )
  
imports: import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import invert
  from tensorflow_probability.python.distributions import exponential
  from tensorflow_probability.python.distributions import independent
  from tensorflow_probability.python.distributions import joint_distribution_coroutine
  from tensorflow_probability.python.distributions import log_prob_ratio
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.distributions import sample
  from tensorflow_probability.python.distributions import transformed_distribution
  from tensorflow_probability.python.experimental.distribute import joint_distribution
  from tensorflow_probability.python.experimental.distribute import sharded
  from tensorflow_probability.python.experimental.distributions import increment_log_prob
  from tensorflow_probability.python.internal import distribute_test_lib
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.math import gradient
label_desc: Use image node_id 5 for calling the ShardTest obj's underlying member method code with example usage: obj.test_kahan_custom_grad() and returns: gradient