<SYSTEM_TASK:>
Solve the following problem using Python, implementing the functions described below, one line at a time
<END_TASK>
<USER_TASK:>
Description:
 def _tf_flatten_batch_dims(x, num_nonbatch_dims):
 """Flatten all but last num_nonbatch_dims into one dimension. Args: x: a tf.Tensor: num_nonbatch_dims: an integer Returns: a tf.Tensor with 1 + num_nonbatch_dims dimensions. """

  shape = x.shape.as_list()
  assert None not in shape
  new_shape = ([list_product(shape[:-num_nonbatch_dims])]
               + shape[-num_nonbatch_dims:])
  if new_shape != shape:
    x = tf.reshape(x, new_shape)
  return x