Fix ONNX exports for Optimum compatible models
Author: merveenoyan
Pull request: yes

@amyeroberts as discussed and also pinging @xenova for review :') (who also fixed DPT) 

I prioritized Optimum compatible ones because I'm launching a project where there's Optimum examples for vision models. I will have a separate PR for the models that aren't compatible with Optimum. Rest of the Optimum compatible models export well without a problem.

Comments:
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_31311). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

@amyeroberts the failing tests seem irrelevant to this PR, I can't re-run them, can you re-run?

@merveenoyan si si - done! 

@amyeroberts to my understanding, torch ONNX export internally calls `trace` thus to capture the graph properly all needs to be torch tensor, one needs to check tracing, and convert to torch tensor/convert ops internally if it's not torch

@amyeroberts can you merge if you think it's ok?

> @amyeroberts to my understanding, torch ONNX export internally calls trace thus to capture the graph properly all needs to be torch tensor, one needs to check tracing, and convert to torch tensor/convert ops internally if it's not torch

Right, I see why we need to do it for the onnx export, but for day-to-day use could we just use torch primitives instead of a python `int` or `float` i.e. do we need to maintain this if/else structure or can we switch everything to torch land regardless of whether we're tracing or not?

@amyeroberts I guess if it's just torch modelling code then yes. Would you like me to swap everything? 

also asking the same question to @xenova 

> Would you like me to swap everything?

@merveenoyan Yes please! This will be cleaner and easier to follow in the code :) 

> also asking the same question to @xenova

I agree with @amyeroberts - if there is a way to "do everything in torch land", that's the best solution! However, there are cases where I'm not entirely sure how to do this. For example, with `torch.nn.interpolate`:
- during inference, the sizes/scales CANNOT be tensors. They must be python ints/floats.
- during tracing, the sizes/scales SHOULD be tensors. Dynamic shapes usually break if they are python ints/floats due to the loss of tracing.

See [here](https://github.com/YvanYin/Metric3D/issues/117#issue-2364404569) for example code (DinoV2 backbone):
```py
        if torch.jit.is_tracing():
          sqrt_N = N ** 0.5
          patch_pos_embed = nn.functional.interpolate(
              patch_pos_embed.reshape(1, (sqrt_N).to(torch.int64), (sqrt_N).to(torch.int64), dim).permute(0, 3, 1, 2),
              size=(w0, h0),
              mode="bicubic",
              antialias=self.interpolate_antialias,
          )
        else:
          sqrt_N = math.sqrt(N)
          sx, sy = float(w0) / sqrt_N, float(h0) / sqrt_N
          patch_pos_embed = nn.functional.interpolate(
              patch_pos_embed.reshape(1, int(sqrt_N), int(sqrt_N), dim).permute(0, 3, 1, 2),
              scale_factor=(sx, sy),
              mode="bicubic",
              antialias=self.interpolate_antialias,
          )
```

Very ugly... I know :/ 



@xenova sounds good, very glad to work with you tbh I didn't know that it would be required in inference.
@amyeroberts should we merge? 

@merveenoyan My understanding from above was that the PR would be updated to remove all the if/else structures wherever possible (but as @xenova points out isn't everywhere unfortunately)

@amyeroberts from what I understood we should still keep them in if/else not to break the inference (I'm also scared of edge cases if there is etc) so I'd rather keep them. what I can do is to test all of them to see if they break or not when all are tensors and remove where it doesn't have to be a python type

@merveenoyan OK. Let's just merge then and we can follow up in future PRs 👍