code: def dimension_size(x, idx):
      """Equivalent to `shape(x)[idx]`, but robust to partially-known shapes."""
      x_shape = x.shape
      static_idx = tf.get_static_value(idx)
      if static_idx is not None:
          idx = static_idx
          if tf.compat.dimension_value(x_shape[idx]) is not None:
              return tf.compat.dimension_value(x_shape[idx])
          x_shape = _shape(x)
      else:  # `dimension_idx` is not static.
          x_shape = convert_to_shape_tensor(_shape(x))
      return x_shape[idx]
  
imports: import decorator
  import numpy
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow_probability.python.internal.backend import numpy
  from tensorflow.python.framework import tensor_util
  from tensorflow.python.ops import control_flow_case
  from tensorflow.python.util import tf_inspect
label_desc: Use image node_id 14 for calling a global function with example usage: dimension_size(x, idx) and returns: x_shape, tf