def maybe_flatten4d3d(x):
  """Flatten if tensor has 4 dimensions.

  Pass through otherwise.

  This is useful since additional dimensions are sometimes removed on the TPU,
  see e.g.
    https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/models/transformer.py?l=159&rcl=279807999

  Args:
    x: a tensor

  Returns:
    A 3D tensor if x is 4D, unmodified x otherwise.
  """
  xshape = common_layers.shape_list(x)
  return common_layers.flatten4d3d(x) if len(xshape) == 4 else x
Flatten if tensor has 4 dimensions.

  Pass through otherwise.

  This is useful since additional dimensions are sometimes removed on the TPU,
  see e.g.
    https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/models/transformer.py?l=159&rcl=279807999

  Args:
    x: a tensor

  Returns:
    A 3D tensor if x is 4D, unmodified x otherwise.
  
tensor2tensor/models/research/transformer_seq2edits.py