refactor almost identical tests
Author: stas00
Pull request: yes

in preparation for adding more schedulers this PR refactors these almost identical tests.

Unfortunately [can't use `pytest.mark.parametrize`](https://docs.pytest.org/en/latest/unittest.html#pytest-features-in-unittest-testcase-subclasses), so the only drawback that it makes them all into a single test. It'd have been nice to parametrize instead.

Comments:
# [Codecov](https://codecov.io/gh/huggingface/transformers/pull/6339?src=pr&el=h1) Report
> Merging [#6339](https://codecov.io/gh/huggingface/transformers/pull/6339?src=pr&el=desc) into [master](https://codecov.io/gh/huggingface/transformers/commit/322dffc6c9a44fd504b24b0efcbcaa419b577a93&el=desc) will **decrease** coverage by `0.03%`.
> The diff coverage is `n/a`.

[![Impacted file tree graph](https://codecov.io/gh/huggingface/transformers/pull/6339/graphs/tree.svg?width=650&height=150&src=pr&token=9qOlN6Hb1c)](https://codecov.io/gh/huggingface/transformers/pull/6339?src=pr&el=tree)

```diff
@@            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](https://codecov.io/gh/huggingface/transformers/pull/6339?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [src/transformers/generation\_tf\_utils.py](https://codecov.io/gh/huggingface/transformers/pull/6339/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9nZW5lcmF0aW9uX3RmX3V0aWxzLnB5) | `84.21% <0.00%> (-2.26%)` | :arrow_down: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/huggingface/transformers/pull/6339?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/huggingface/transformers/pull/6339?src=pr&el=footer). Last update [322dffc...92d3825](https://codecov.io/gh/huggingface/transformers/pull/6339?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).


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.

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

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>
```

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: https://github.com/huggingface/transformers/blob/175cd45e13b2e33d1efec9e2ac217cba99f6ae58/examples/seq2seq/test_seq2seq_examples.py#L238


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