def framework(msg='err'):
  """Return framework module or dummy version."""
  del msg
  if is_tf2:
    return DummyModule(
        arg_scope=None,
        get_name_scope=lambda: tf.get_default_graph().get_name_scope(),
        name_scope=tf.name_scope,
        deprecated=deprecated,
        nest=tf.nest,
        argsort=tf.argsort)

  from tensorflow.contrib import framework as contrib_framework  # pylint: disable=g-direct-tensorflow-import,g-import-not-at-top
  return contrib_framework
Return framework module or dummy version.
tensor2tensor/utils/contrib.py