def _axis_size(x, axis=None):
  """"""
  if axis is None:
    return tf.cast(tf.size(input=x), x.dtype)
  return tf.cast(
      tf.reduce_prod(input_tensor=tf.gather(tf.shape(input=x), axis)), x.dtype)