[Reformer] Simpler reverse sort backward implementation
Author: patrickvonplaten
Pull request: yes

After reviewing code of similar PyTorch implementation: https://github.com/lucidrains/reformer-pytorch/pull/104 and original trax code again: https://github.com/google/trax/blob/master/trax/layers/research/efficient_attention.py#L1265-L1266, there is a simpler way to implement the backward reverse sort function. 
The code is taken from https://github.com/lucidrains/reformer-pytorch/blob/42a8682ff8e7cec3122eff6febc9087f1c53f370/reformer_pytorch/reformer_pytorch.py#L414.

This simpler code is also implemented in the Reformer test branch and all tests are checked for correctness: https://github.com/huggingface/transformers/tree/branch_to_save_trax_integration_tests


**Note**: There was no bug in the code before / the logic of the code has not changed. This PR just makes the backward function of ReverseSort simpler.


Comments:
# [Codecov](https://codecov.io/gh/huggingface/transformers/pull/5343?src=pr&el=h1) Report
> Merging [#5343](https://codecov.io/gh/huggingface/transformers/pull/5343?src=pr&el=desc) into [master](https://codecov.io/gh/huggingface/transformers/commit/9fe09cec76efa1e221c3fd6eb8520ba0a911f092&el=desc) will **increase** coverage by `0.40%`.
> The diff coverage is `66.66%`.

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

```diff
@@            Coverage Diff             @@
##           master    #5343      +/-   ##
==========================================
+ Coverage   77.93%   78.33%   +0.40%     
==========================================
  Files         138      138              
  Lines       23860    23851       -9     
==========================================
+ Hits        18595    18684      +89     
+ Misses       5265     5167      -98     
```


| [Impacted Files](https://codecov.io/gh/huggingface/transformers/pull/5343?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [src/transformers/modeling\_reformer.py](https://codecov.io/gh/huggingface/transformers/pull/5343/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ19yZWZvcm1lci5weQ==) | `89.18% <66.66%> (+0.97%)` | :arrow_up: |
| [src/transformers/modeling\_tf\_utils.py](https://codecov.io/gh/huggingface/transformers/pull/5343/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ190Zl91dGlscy5weQ==) | `85.41% <0.00%> (-0.74%)` | :arrow_down: |
| [src/transformers/tokenization\_utils.py](https://codecov.io/gh/huggingface/transformers/pull/5343/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy90b2tlbml6YXRpb25fdXRpbHMucHk=) | `91.82% <0.00%> (+0.31%)` | :arrow_up: |
| [src/transformers/tokenization\_bert.py](https://codecov.io/gh/huggingface/transformers/pull/5343/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy90b2tlbml6YXRpb25fYmVydC5weQ==) | `90.86% <0.00%> (+0.91%)` | :arrow_up: |
| [src/transformers/modeling\_openai.py](https://codecov.io/gh/huggingface/transformers/pull/5343/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ19vcGVuYWkucHk=) | `80.90% <0.00%> (+1.38%)` | :arrow_up: |
| [src/transformers/data/processors/squad.py](https://codecov.io/gh/huggingface/transformers/pull/5343/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9kYXRhL3Byb2Nlc3NvcnMvc3F1YWQucHk=) | `56.68% <0.00%> (+28.02%)` | :arrow_up: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/huggingface/transformers/pull/5343?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/5343?src=pr&el=footer). Last update [9fe09ce...d53f899](https://codecov.io/gh/huggingface/transformers/pull/5343?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).