code: def test_dimension_size(self):
      shape = [3, 2, 1]
      idx = 1
      idx_tensor = tf.convert_to_tensor(idx)
  
      # case: numpy input.
      self.assertEqual(
          ps.dimension_size(np.zeros(shape), idx), shape[idx]
      )
      self.assertEqual(
          ps.dimension_size(np.zeros(shape), idx_tensor), shape[idx]
      )
  
      # case: static-shape Tensor input.
      self.assertEqual(
          ps.dimension_size(tf.zeros(shape), idx), shape[idx]
      )
      self.assertNotIsInstance(
          ps.dimension_size(tf.zeros(shape), idx), tf1.Dimension
      )
      self.assertEqual(
          ps.dimension_size(tf.zeros(shape), idx_tensor), shape[idx]
      )
      self.assertNotIsInstance(
          ps.dimension_size(tf.zeros(shape), idx_tensor), tf1.Dimension
      )
  
      if tf.executing_eagerly():
          return
  
      # Case: input is Tensor with fully unknown shape.
      zeros_pl = tf1.placeholder_with_default(
          tf.zeros(shape), shape=None
      )
      idx_pl = tf1.placeholder_with_default(idx, shape=[])
      self.assertAllEqual(ps.dimension_size(zeros_pl, idx), shape[idx])
      self.assertAllEqual(
          ps.dimension_size(zeros_pl, idx_tensor), shape[idx]
      )
      self.assertAllEqual(
          ps.dimension_size(zeros_pl, idx_pl), shape[idx]
      )
  
      # Case: input is Tensor with partially known shape.
      # The result should be static if idx is.
      zeros_partial_pl = tf1.placeholder_with_default(
          tf.zeros(shape),
          shape=tf.TensorShape([None, shape[idx], None]),
      )
      self.assertEqual(
          ps.dimension_size(zeros_partial_pl, idx), shape[idx]
      )
      self.assertEqual(
          ps.dimension_size(zeros_partial_pl, idx_tensor), shape[idx]
      )
      self.assertAllEqual(
          ps.dimension_size(zeros_partial_pl, idx_pl), shape[idx]
      )
  
imports: from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import test_util
label_desc: Use image node_id 4 for calling the ShapeTest obj's underlying member method code with example usage: obj.test_dimension_size() without return types