code: def testPartialShape(self):
      @tf.function
      def f(x):
          model_dist = jds.JointDistributionSequential(
              [sample_dist_lib.Sample(normal.Normal(0, 1), tf.shape(x))]
          )
          return initialization.init_near_unconstrained_zero(
              model_dist
          ).sample(seed=test_util.test_seed())
  
      g = f.get_concrete_function(tf.TensorSpec([None, 3], tf.float32))
      self.evaluate(g(tf.ones([7, 3])))
  
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 6 for calling the InitializationTest obj's underlying member method code with example usage: obj.testPartialShape() and returns: initialization