import tensorflow as tf



# TODO(zoyahav): Use register_keras_serializable directly once we no longer support
# TF<2.1.
def _maybe_register_keras_serializable(package):
  if hasattr(tf.keras.utils, 'register_keras_serializable'):
    return tf.keras.utils.register_keras_serializable(package=package)
  else:
    return lambda cls: cls


def _check_tensorflow_version():
