code: def test_can_initialize_from_sharded_distribution(self):
      def model():
          x = yield normal.Normal(0.0, 1.0)
          yield distribute.Sharded(
              normal.Normal(x, 1.0), self.axis_name
          )
  
      jd = distribute.JointDistributionCoroutine(model)
  
      def run(seed):
          init_jd = initialization.init_near_unconstrained_zero(jd)
          return init_jd.sample(seed=seed)
  
      x, y = self.evaluate(
          self.per_replica_to_tensor(
              self.strategy_run(
                  run, args=(test_util.test_seed(),), in_axes=None
              )
          )
      )
      for i in range(distribute_test_lib.NUM_DEVICES):
          for j in range(distribute_test_lib.NUM_DEVICES):
              if i == j:
                  continue
              self.assertAllClose(x[i], x[j])
              self.assertNotAllClose(y[i], y[j])
  
imports: import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import exp
  from tensorflow_probability.python.bijectors import identity
  from tensorflow_probability.python.bijectors import pad
  from tensorflow_probability.python.bijectors import sigmoid
  from tensorflow_probability.python.distributions import half_cauchy
  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 multinomial
  from tensorflow_probability.python.distributions import mvn_diag
  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 import distribute
  from tensorflow_probability.python.experimental.mcmc import initialization
  from tensorflow_probability.python.internal import distribute_test_lib
  from tensorflow_probability.python.internal import test_util
label_desc: Use image node_id 1 for calling the DistributedTest obj's underlying member method code with example usage: obj.test_can_initialize_from_sharded_distribution() and returns: init_jd