code: def _squeeze(x, axis):
      """A version of squeeze that works with dynamic axis."""
      x = tf.convert_to_tensor(x, name="x")
      if axis is None:
          return tf.squeeze(x, axis=None)
      axis = ps.convert_to_shape_tensor(
          axis, name="axis", dtype=tf.int32
      )
      axis = _make_list_or_1d_tensor(axis)  # Ensure at least 1d.
      keep_axis = ps.setdiff1d(ps.range(0, ps.rank(x)), axis)
      return tf.reshape(x, ps.gather(ps.shape(x), keep_axis))
  
imports: import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import assert_util
  from tensorflow_probability.python.internal import distribution_util
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow_probability.python.math.generic import reduce_logmeanexp
label_desc: Use image node_id 20 for calling a global function with example usage: _squeeze(x, axis) and returns: tf, tf