try:
        # Dynamically try to load the tf v2 lookup, falling back to contrib
        lookup = tf.compat.v2.lookup
        hash_table_class = tf.compat.v2.lookup.StaticHashTable
      except AttributeError: