MaskFormer,Mask2former - reduce memory load
Author: amyeroberts
Pull request: yes

# What does this PR do?

#25297 resolved an error that occurred when trying to trace MaskFormer-like models. The solution was to remove einsum operations. However, the replacements were very memory intensive: a large matrix was created and then summed over. 

This PR replaces this logic by creating the result array and the iterates over the summed dimension. 

Fixes #25709

It was confirmed that this resolved the memory issue, by installing from this branch and rerunning the example notebook successfully: https://colab.research.google.com/drive/1xq54l9a2AQLIHT5jw63btifbOrvSqzts?usp=sharing


## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#start-contributing-pull-requests),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?

Comments:
_The documentation is not available anymore as the PR was closed or merged._

@ydshieh, I ran the following to check the times. The einsum operation is quicker, but on the same order of magnitude as this and they scale the same way. This I expected as both will still need to reduce over the channel dimensions and this should (and does) scale linearly. 


```python
def foo(mask_embeddings, pixel_embeddings):
    batch_size, num_queries, num_channels = mask_embeddings.shape
    _, _, height, width = pixel_embeddings.shape
    binaries_masks = np.zeros((batch_size, num_queries, height, width))
    for c in range(num_channels):
        binaries_masks += mask_embeddings[..., c][..., None, None] * pixel_embeddings[:, None, c]
    return binaries_masks

def bar(mask_embeddings, pixel_embeddings):
    binaries_masks = np.einsum('bqc, bchw -> bqhw', mask_embeddings, pixel_embeddings)
    return binaries_masks

foo_times = {}
bar_times = {}

for num_channels in (10, 100, 1000, 10000):
    foo_times[num_channels] = []
    bar_times[num_channels] = []

    mask_embeddings = np.random.randn(batch_size, num_queries, num_channels)
    pixel_embeddings = np.random.randn(batch_size, num_channels, height, width)

    for i in range(100):
        start = time.time()
        foo(mask_embeddings, pixel_embeddings)
        foo_times[num_channels].append(time.time() - start)
        start = time.time()
        bar(mask_embeddings, pixel_embeddings)
        bar_times[num_channels].append(time.time() - start)

    print("\n", num_channels)
    print(f"Iterate directly: {np.mean(foo_times[num_channels]):.5f} +/- {np.std(foo_times[num_channels]):.5f}")
    print(f"Einsum approach:  {np.mean(bar_times[num_channels]):.5f} +/- {np.std(bar_times[num_channels]):.5f}")
```

Results: 
```
 10
Iterate directly: 0.00038 +/- 0.00001
Einsum approach:  0.00021 +/- 0.00000

 100
Iterate directly: 0.00300 +/- 0.00008
Einsum approach:  0.00203 +/- 0.00004

 1000
Iterate directly: 0.03329 +/- 0.00320
Einsum approach:  0.02006 +/- 0.00041

 10000
Iterate directly: 0.35071 +/- 0.02558
Einsum approach:  0.20384 +/- 0.00132
```

The question is: are we happy with this step being almost 2x slower if it enables jit tracing? As the original issue arises from a torch operation - torch.einsum - arguably it's not something we should be addressing on our end. 

I do have a slight concern with the approach in this PR wrt numerical stability. 



Yeah, not ideal although not too bad. If the 2x slower is OK depends on its scale: 20 v.s 10 seconds is not good, and 0.02 v.s 0.01 is kind fine. I don't know what's the value for `num_queries` when we use a real checkpoint (and what value you used above) and therefore what is the usual op timing scale here.

Are you up to use 

```python3
if config.torchscript:
    iterate way
else:
   einsum
```
?

If so, let's do it and we don't have to worry about the timing. Otherwise, let's go this PR but have a comment about the potential slower op time, a link to your above comment and a link to the original enisum implementation.


@ydshieh The number of channels in on the order of 100s. For the checkpoint `""facebook/maskformer-swin-base-ade"` it's 256. This means the time difference between the two approaches is ~0.001 s - 0.01s.

I'm not going to add the if/else for the moment for the sake of simplicity. Looking in the library, I don't see this control flow being used often and I'm not confident that it'll be correctly set when our users want to trace the model. 

> @ydshieh The number of channels in on the order of 100s. For the checkpoint `""facebook/maskformer-swin-base-ade"` it's 256. This means the time difference between the two approaches is ~0.001 s - 0.01s.
> 
> I'm not going to add the if/else for the moment for the sake of simplicity. Looking in the library, I don't see this control flow being used often and I'm not confident that it'll be correctly set when our users want to trace the model.

Sure, works for me, let's keep it simple.

Regarding set `torchscript`, I  think it is required, see `return_dict = return_dict if return_dict is not None else self.config.return_dict`. Also it appears in the original issue description #25261.