from tensorflow.python import tf2



def _check_tensorflow_version():
  """Check that we're using a compatible TF version.

  Raises a warning if either Tensorflow version is less that 2.0 or TF 2.x is
  not enabled.

  If TF 2.x is enabled, but version is < TF 2.3, raises a warning to indicate
  that resources may not be initialized.
  """
  major, minor, _ = tf.version.VERSION.split('.')
  if not (int(major) >= 2 and tf2.enabled()):
    tf.compat.v1.logging.warning(
        'Tensorflow version (%s) found. TransformFeaturesLayer is supported '
        'only for TF 2.x with TF 2.x behaviors enabled and may not work as '
        'intended.', tf.version.VERSION)
  elif int(major) == 2 and int(minor) < 3:
    # TODO(varshaan): Log a more specific warning.
    tf.compat.v1.logging.warning(
        'Tensorflow version (%s) found. TransformFeaturesLayer may not work '
        'as intended if the SavedModel contains an initialization op.',
        tf.version.VERSION)

