def maybe_flatten3d2d(x):
  """Flatten if tensor has 3 dimensions, similar to maybe_flatten4d3d()."""
  xshape = common_layers.shape_list(x)
  if len(xshape) != 3:
    return x
  return tf.reshape(x, [xshape[0], xshape[1] * xshape[2]])
Flatten if tensor has 3 dimensions, similar to maybe_flatten4d3d().
tensor2tensor/models/research/transformer_seq2edits.py