tf.app.flags.DEFINE_integer('num_blocks', 1, 'Number of blocks in each attention')
tf.app.flags.DEFINE_integer('num_heads', 8, 'Number of heads in each attention')

FLAGS = tf.app.flags.FLAGS
__flags = FLAGS.__flags if tf.__version__ < '1.15.0' else FLAGS.flag_values_dict()