code: def flatten_all_but_last(a):
      """Flatten all dimensions of a except the last."""
      ret = tf.reshape(a, [-1, tf.shape(a)[-1]])
      if not tf.executing_eagerly():
          ret.set_shape([None] + a.get_shape().as_list()[-1:])
      return ret
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import functools
  import math
  import six
  from six.moves import range
  from six.moves import zip
  from tensor2tensor.layers import common_layers
  from tensor2tensor.layers.vq_discrete import DiscreteBottleneck
  import tensorflow.compat.v1
label_desc: Use image node_id 17 for calling a global function with example usage: flatten_all_but_last(a) and returns: ret