refactor almost identical tests

[stas00]

in preparation for adding more schedulers this PR refactors these almost identical tests.

Unfortunately can't use pytest.mark.parametrize, so the only drawback that it makes them all into a single test. It'd have been nice to parametrize instead.

[codecov]

Codecov Report

Merging #6339 into master will decrease coverage by 0.03%.
The diff coverage is n/a.

@@            Coverage Diff             @@
##           master    #6339      +/-   ##
==========================================
- Coverage   78.37%   78.34%   -0.04%     
==========================================
  Files         148      148              
  Lines       27196    27196              
==========================================
- Hits        21316    21307       -9     
- Misses       5880     5889       +9     

Impacted Files	Coverage Δ
src/transformers/generation_tf_utils.py	84.21% <0.00%> (-2.26%)	⬇️

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 322dffc...92d3825. Read the comment docs.

[stas00]

could also modify unwrap_schedule and unwrap_and_save_reload_schedule to return a clean list of numbers, and then it'd be just:

        for scheduler_func, data in scheds.items():
            kwargs, expected_learning_rates = data

            scheduler = scheduler_func(self.optimizer, **kwargs)
            lrs_1 = unwrap_schedule(scheduler, self.num_steps)
            self.assertListAlmostEqual(lrs_1, expected_learning_rates, tol=1e-2)

            scheduler = scheduler_func(self.optimizer, **kwargs)
            lrs_2 = unwrap_and_save_reload_schedule(scheduler, self.num_steps)
            self.assertListEqual(lrs_1, lrs_2)

but perhaps it'd be less intuitive for those reading the test code.

[sshleifer]

Does this impact tracebacks in a bad way? Previously I would know which scheduler I broke if test_warmup_constant_scheduler failed.

[stas00]

That's super-imporant, @sshleifer, thank you for flagging that!

Added an assert msg to make it clear what fails, e.g. if I break data for the sake of demo, we now get:

        for scheduler_func, data in scheds.items():
            kwargs, expected_learning_rates = data

            scheduler = scheduler_func(self.optimizer, **kwargs)
            lrs_1 = unwrap_schedule(scheduler, self.num_steps)
            self.assertEqual(len(lrs_1[0]), 1)
            self.assertListAlmostEqual(
>               [l[0] for l in lrs_1], expected_learning_rates, tol=1e-2, msg=f"failed for {scheduler_func}"
            )

tests/test_optimization.py:126:
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
tests/test_optimization.py:92: in assertListAlmostEqual
    self.assertAlmostEqual(a, b, delta=tol, msg=msg)
E   AssertionError: 2.5 != 3.5 within 0.01 delta (1.0 difference) : failed for <function get_constant_schedule_with_warmup at 0x7f5da6f0bdd0>

[stas00]

hmm, not sure whether the last commit, to make the assert message even more specific, was needed.

Also, alternatively, I can move the code out of unittest class and then use pytest parametrization so it'll be self-documenting on assert. Ala:

transformers/examples/seq2seq/test_seq2seq_examples.py

Line 238 in 175cd45

@pytest.mark.parametrize(

[sshleifer]

LGTM as is, but won't merge it myself.

[LysandreJik]

Looks good to me, looking forward to new schedulers!