ONNX export may fail for Hiera due to math.sqrt and python type int casts

[xenova]

Similar to a previous issue (#31311) we had with vision models. cc @merveenoyan @amyeroberts @NielsRogge

Offending lines:

transformers/src/transformers/models/hiera/modeling_hiera.py

Line 340 in 5c1027b

pos_embeds = pos_embeds.reshape(1, int(math.sqrt(num_positions)), int(math.sqrt(num_positions)), dim)

transformers/src/transformers/models/hiera/modeling_hiera.py

Line 344 in 5c1027b

scale_factor=(h0 / math.sqrt(num_positions), w0 / math.sqrt(num_positions)),

We should probably abstract interpolate_pos_encoding, since this is reused across many different vision architectures.

[qubvel]

Hi @xenova, there are more issues with the current implementation of interpolate_pos_encoding

• Potential error in num_patches calculation in src/transformers/models/vit_mae/modeling_vit_mae.py #32410 (comment)

Maybe it's better to rewrite interpolation using size argument instead of scale to avoid errors described in the issue above, will it be compatible with ONNX export? WDYT?

[xenova]

Maybe it's better to rewrite interpolation using size argument instead of scale to avoid errors described in the issue above, will it be compatible with ONNX export? WDYT?

Yes I agree! Some other issues/PRs where this is discussed:

• Improved ONNX support with dynamic shapes YvanYin/Metric3D#117 (comment)
• Add ONNX export for DinoV2 models optimum#1580
• Make vitdet jit trace complient #30065

I'll try draft a PR for this soon.