code: def make_jd_coroutine(axis_name):
      def model_coroutine():
          w = yield Root(normal.Normal(0.0, 1.0))
          x = yield sharded.Sharded(
              sample_dist_lib.Sample(normal.Normal(w, 1.0), 1),
              shard_axis_name=axis_name,
          )
          yield sharded.Sharded(
              independent.Independent(normal.Normal(x, 1.0), 1),
              shard_axis_name=axis_name,
          )
  
      return jd.JointDistributionCoroutine(model_coroutine)
  
imports: from absl.testing import parameterized
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import scale
  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 lognormal
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.distributions import sample
  from tensorflow_probability.python.distributions import uniform
  from tensorflow_probability.python.experimental.distribute import joint_distribution
  from tensorflow_probability.python.experimental.distribute import sharded
  from tensorflow_probability.python.internal import distribute_test_lib
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.math import gradient
label_desc: Use image node_id 4 for calling a global function with example usage: make_jd_coroutine(axis_name) and returns: jd