def layers():
  """Get the layers module good for TF 1 and TF 2 work for now."""
  layers_module = None
  try:
    layers_module = tf.layers
  except AttributeError:
    logging.info("Cannot access tf.layers, trying TF2 layers.")
  try:
    from tensorflow.python import tf2  # pylint: disable=g-direct-tensorflow-import,g-import-not-at-top
    if tf2.enabled():
      logging.info("Running in V2 mode, using Keras layers.")
      layers_module = tf.keras.layers
  except ImportError:
    pass
  return layers_module
Get the layers module good for TF 1 and TF 2 work for now.
tensor2tensor/layers/common_layers.py