code: def test_kahan_precision(self, jd_class, jit):
      maybe_jit = lambda f: f
      if jit:
          self.skip_if_no_xla()
          if not JAX_MODE and not tf.test.is_gpu_available():
              self.skipTest("b/179303849")
          maybe_jit = tf.function(jit_compile=True)
  
      def make_models(dtype):
          models = {}
  
          def mk_20k_poisson(log_rate):
              return poisson.Poisson(
                  log_rate=tf.broadcast_to(
                      log_rate[..., tf.newaxis],
                      log_rate.shape + (20_000,),
                  )
              )
  
          def coroutine_model():
              log_rate = yield normal.Normal(
                  0.0, dtype(0.2), name="log_rate"
              )
              yield mk_20k_poisson(log_rate).copy(name="x")
  
          models[
              jdab.JointDistributionCoroutineAutoBatched
          ] = coroutine_model
  
          models[jdab.JointDistributionSequentialAutoBatched] = [
              normal.Normal(0.0, dtype(0.2)),
              mk_20k_poisson,
          ]
  
          models[
              jdab.JointDistributionNamedAutoBatched
          ] = collections.OrderedDict(
              (
                  ("log_rate", normal.Normal(0.0, dtype(0.2))),
                  ("x", mk_20k_poisson),
              )
          )
          return models
  
      joint = jd_class(
          make_models(np.float32)[jd_class],
          validate_args=True,
          experimental_use_kahan_sum=True,
      )
      joint64 = jd_class(
          make_models(np.float64)[jd_class], validate_args=True
      )
      stream = test_util.test_seed_stream()
      nsamp = 7
      xs = self.evaluate(
          joint.sample(log_rate=tf.zeros([nsamp]), seed=stream())
      )
      if isinstance(xs, dict):
          xs["log_rate"] = normal.Normal(0, 0.2).sample(
              nsamp, seed=stream()
          )
      else:
          xs = (
              normal.Normal(0, 0.2).sample(nsamp, seed=stream()),
              xs[1],
          )
      xs64 = tf.nest.map_structure(lambda x: tf.cast(x, tf.float64), xs)
      lp = maybe_jit(joint.copy(validate_args=not jit).log_prob)(xs)
      lp64 = joint64.log_prob(xs64)
      lp, lp64 = self.evaluate((tf.cast(lp, tf.float64), lp64))
      # Without Kahan, example max-abs-diff: ~0.06
      self.assertAllClose(lp64, lp, rtol=0.0, atol=0.01)
  
imports: import collections
  import os
  from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import bijector_test_util
  from tensorflow_probability.python.bijectors import exp
  from tensorflow_probability.python.bijectors import softplus
  from tensorflow_probability.python.distributions import bernoulli
  from tensorflow_probability.python.distributions import deterministic
  from tensorflow_probability.python.distributions import dirichlet
  from tensorflow_probability.python.distributions import exponential
  from tensorflow_probability.python.distributions import half_normal
  from tensorflow_probability.python.distributions import inverse_gamma
  from tensorflow_probability.python.distributions import joint_distribution_auto_batched
  from tensorflow_probability.python.distributions import joint_distribution_coroutine
  from tensorflow_probability.python.distributions import joint_distribution_named
  from tensorflow_probability.python.distributions import joint_distribution_sequential
  from tensorflow_probability.python.distributions import lognormal
  from tensorflow_probability.python.distributions import multinomial
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.distributions import poisson
  from tensorflow_probability.python.distributions import sample
  from tensorflow_probability.python.distributions import student_t
  from tensorflow_probability.python.distributions import transformed_distribution
  from tensorflow_probability.python.distributions import uniform
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.util import deferred_tensor
label_desc: Use image node_id 22 for calling the JointDistributionAutoBatchedTest obj's underlying member method code with example usage: obj.test_kahan_precision(jd_class, jit) and returns: models, poisson