import tensorflow as tf

        step_size = tf.cast(step_size, dtype=tf.float32)
        max_lr = tf.cast(max_lr, dtype=tf.float32)

        if mode == 'tri':
            periodic_comp = tf.mod((global_step + step_size / 4) / step_size, 1)
            first_factor = tf.abs(periodic_comp - 0.5)
            second_factor = 2 * (max_lr - learning_rate)
            second_comp = learning_rate
