fix(mps): build_2d_sinusoidal_position_embedding crashes on Apple Silicon due to float64 on MPS device
Author: shubhammr21
Pull request: yes

Fixes #46159

## Problem

`build_2d_sinusoidal_position_embedding` hardcodes `torch.float64` **and** passes `device=device` to every intermediate tensor. MPS (Apple Silicon) does not support float64, so any model calling this function with an MPS device crashes with:

```
TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework
doesn't support float64. Please use float32 instead.
```

## Why float64 was there

The original design was intentional: `1.0 / temperature**omega` (temperature=10000) produces very small values where float32 loses significant digits. The function computed everything at float64 precision then cast to the output dtype at the end via `.to(dtype)`. The docstring documented this: *"frequency arithmetic uses float64 internally"*.

Simply replacing `float64` with `dtype` throughout would fix MPS but silently degrade precision for `float16`/`bfloat16` callers.

## Fix

Keep float64 for intermediate arithmetic, but on **CPU** (where float64 always works). Only the final result is moved to the requested device and cast to dtype:

```python
# Before — float64 on the target device → crashes on MPS
omega = torch.arange(pos_dim, dtype=torch.float64, device=device) / pos_dim
...
return pos_embed.to(dtype)

# After — float64 on CPU, move to device at the end → works everywhere
omega = torch.arange(pos_dim, dtype=torch.float64) / pos_dim  # no device=
...
return pos_embed.to(device=device, dtype=dtype)
```

## Scope

The canonical fix is one function in `modular_vit_mae.py`. `make fix-repo` propagated it to all 8 generated files that inline this function: `modeling_vit_mae`, `modeling_rt_detr`, `modeling_rt_detr_v2`, `modeling_aimv2`, `modeling_deimv2`, `modeling_d_fine`, `modeling_pp_doclayout_v2`, `modeling_pp_doclayout_v3`.

## Testing

| Platform | Command | Result |
|---|---|---|
| macOS MPS — Apple Silicon | smoke test: float32 / float16 / bfloat16 | ✅ all pass |
| Linux CPU — arm64 Docker (`python:3.12-slim` + torch CPU wheel) | smoke test: float32 / float16 / bfloat16 / precision check | ✅ all pass |
| Unit tests `vit_mae`, `rt_detr`, `rt_detr_v2` | `pytest tests/models/vit_mae/ tests/models/rt_detr/ tests/models/rt_detr_v2/ -q` | ✅ 288 passed, 0 failed |
| Repo checks | `make fix-repo` | ✅ all 15 checks pass |

> **AI assistance disclosure:** Claude was used to identify the affected modular files and run the fix-repo workflow. Root cause analysis, fix approach, and MPS verification were done by the issue author on real Apple Silicon hardware.

Comments:
**[For maintainers]** Suggested jobs to run (before merge)

run-slow: aimv2, d_fine, deimv2, pp_doclayout_v2, pp_doclayout_v3, rt_detr, rt_detr_v2, vit_mae

Thanks for the review! On the bit-identical question — confirmed, precision test showed `max_err = 0.00e+00` between float64 reference and float32 output.

One honest gap: I only have an MPS device to test on, so CUDA wasn't verified locally. The slow CI tests should cover that. Happy to address anything that comes up there.

On the inline comment — will wait for maintainer feedback before pushing another commit.

Hmn, there's an outer multiplication in this path, which we're now doing on CPU. How much does this hurt performance on CUDA systems?

@Rocketknight1 @shubhammr21 — regarding the CUDA performance concern on the `outer()` now running on CPU:

You could scope the CPU fallback to MPS only, leaving the CUDA path fully on-device:

```python
# MPS doesn't support float64 — compute on CPU only for that backend
compute_device = None if device is not None and torch.device(device).type == "mps" else device

pos_dim = embed_dim // 4
omega = torch.arange(pos_dim, dtype=torch.float64, device=compute_device) / pos_dim
omega = 1.0 / temperature**omega  # (D/4,)

grid_h = torch.arange(height, dtype=torch.float64, device=compute_device)
grid_w = torch.arange(width, dtype=torch.float64, device=compute_device)
grid_h, grid_w = torch.meshgrid(grid_h, grid_w, indexing="ij")  # (H, W) each

emb_h = grid_h.flatten().outer(omega)  # (H*W, D/4)
emb_w = grid_w.flatten().outer(omega)  # (H*W, D/4)

pos_embed = torch.cat([emb_h.sin(), emb_h.cos(), emb_w.sin(), emb_w.cos()], dim=1)

if cls_token:
    pos_embed = torch.cat([torch.zeros(1, embed_dim, dtype=torch.float64, device=compute_device), pos_embed], dim=0)

return pos_embed.to(device=device, dtype=dtype)
```

This keeps embeddings numerically consistent with the CUDA/CPU path across all devices (full float64 intermediates everywhere), and answers the perf question by leaving CUDA entirely on-device. The only overhead is a single CPU→MPS transfer for Apple Silicon users, at model init (not in forward).

Verified on MPS (PyTorch 2.12.0, M-series) — `.to(device, dtype)` casts before placing on device, no float64 instantiation on MPS:

```python
import torch
# float64 on CPU → float32 on MPS in one call (safe)
t = torch.randn(10, 256, dtype=torch.float64)
r = t.to(device="mps", dtype=torch.float32)
print(r.dtype, r.device)  # torch.float32 mps:0

# float64 directly on MPS (the original crash)
torch.arange(10, dtype=torch.float64, device="mps")  # raises: MPS doesn't support float64
```

The closed #46513 used a similar device check but downgraded to float32 intermediates on MPS; this keeps full float64 intermediates so the MPS result matches the CPU path.

For `modular_vit_mae.py` it's the same one-line addition (`compute_device = ...`) — propagates cleanly via `make fix-repo`.


I had a look whether we could use fp32 instead of fp64 directly when creating the positional embeddings.

### TLDR
**Recommendation**: Create embeddings in fp64 if device supports it. Create MPS embeddings in fp32. There is no benefit to init with fp64 on CPU and then move to MPS for inference as CPU/MPS inference differences outweigh the impact of f32/fp64 difference. The only motivation I could see for fp64 on CPU init is to create the exact same weights as on other devices. 

---

For normal cpu and cuda usage the differences are negligible. Aimv2 has an 3e-02 `max_abs` difference but this is because it generates high activations (unnormalized in `last_hidden_state`) and `max_rel` stays <1e-3 which means it is fine.

The problem is for users that have TF32 enabled. There the switch from fp64 to fp32 leads to significant output changes (table 3). So my recommendation is to keep fp64 for devices that support it.

Benchmarks ran on an H100.

```
========== SUMMARY (cpu) ==========
           model                  tensor     max_abs    mean_abs     max_rel       cosine_sim
---------------------------------------------------------------------------------------------
         vit_mae   reconstruction_logits   1.550e-06   2.048e-07   6.003e-07   1.000000000000
           aimv2       last_hidden_state   2.784e-02   5.287e-06   4.158e-04   0.999999999655
          d_fine     encoder_feature_map   5.722e-06   1.744e-07   2.931e-07   1.000000000000
          d_fine                  logits   2.521e-03   1.376e-05   2.781e-04   0.999999999958
          d_fine              pred_boxes   1.979e-04   1.543e-06   6.596e-05   0.999999999938
          deimv2     encoder_feature_map   8.345e-06   8.390e-08   8.052e-07   1.000000000000
          deimv2                  logits   3.333e-04   2.761e-05   3.197e-05   0.999999999980
          deimv2              pred_boxes   2.027e-05   6.129e-07   1.691e-05   0.999999999995
 pp_doclayout_v2     encoder_feature_map   2.289e-05   2.052e-07   2.005e-07   1.000000000000
 pp_doclayout_v2                  logits   1.397e-04   2.169e-06   1.662e-05   0.999999999999
 pp_doclayout_v2              pred_boxes   2.384e-06   6.698e-08   2.384e-06   1.000000000000
 pp_doclayout_v3     encoder_feature_map   4.578e-05   2.728e-07   1.799e-07   1.000000000000
 pp_doclayout_v3                  logits   5.150e-05   2.842e-06   5.146e-06   1.000000000000
 pp_doclayout_v3              pred_boxes   7.451e-06   2.083e-07   7.451e-06   1.000000000000
         rt_detr     encoder_feature_map   1.144e-05   1.991e-07   2.848e-07   1.000000000000
         rt_detr                  logits   7.153e-05   2.373e-06   8.649e-06   1.000000000000
         rt_detr              pred_boxes   7.868e-06   1.420e-07   7.868e-06   1.000000000000
      rt_detr_v2     encoder_feature_map   1.907e-05   2.189e-07   3.081e-07   1.000000000000
      rt_detr_v2                  logits   8.440e-05   1.393e-06   8.880e-06   1.000000000000
      rt_detr_v2              pred_boxes   7.808e-06   6.140e-08   7.808e-06   1.000000000000


========== SUMMARY (cuda) ==========
           model                  tensor     max_abs    mean_abs     max_rel       cosine_sim
---------------------------------------------------------------------------------------------
         vit_mae   reconstruction_logits   2.623e-06   2.824e-07   1.023e-06   1.000000000000
           aimv2       last_hidden_state   3.162e-02   7.996e-06   4.722e-04   0.999999999527
          d_fine     encoder_feature_map   1.001e-05   2.543e-07   5.130e-07   1.000000000000
          d_fine                  logits   4.762e-03   3.514e-05   5.252e-04   0.999999999849
          d_fine              pred_boxes   4.845e-04   1.755e-06   1.615e-04   0.999999999778
          deimv2     encoder_feature_map   1.299e-05   8.871e-08   1.254e-06   1.000000000000
          deimv2                  logits   2.251e-04   1.885e-05   2.159e-05   0.999999999993
          deimv2              pred_boxes   1.973e-05   4.252e-07   1.646e-05   0.999999999997
 pp_doclayout_v2     encoder_feature_map   1.526e-05   2.401e-07   1.337e-07   1.000000000000
 pp_doclayout_v2                  logits   2.069e-04   2.912e-06   2.446e-05   0.999999999999
 pp_doclayout_v2              pred_boxes   4.411e-06   8.365e-08   4.411e-06   1.000000000000
 pp_doclayout_v3     encoder_feature_map   5.341e-05   3.138e-07   2.099e-07   1.000000000000
 pp_doclayout_v3                  logits   4.387e-05   2.481e-06   4.429e-06   1.000000000000
 pp_doclayout_v3              pred_boxes   8.285e-06   2.468e-07   8.285e-06   0.999999999999
         rt_detr     encoder_feature_map   9.060e-06   2.287e-07   2.255e-07   1.000000000000
         rt_detr                  logits   6.008e-05   2.209e-06   7.265e-06   1.000000000000
         rt_detr              pred_boxes   8.881e-06   1.593e-07   8.881e-06   1.000000000000
      rt_detr_v2     encoder_feature_map   2.289e-05   2.435e-07   3.698e-07   1.000000000000
      rt_detr_v2                  logits   8.392e-05   1.437e-06   8.830e-06   1.000000000000
      rt_detr_v2              pred_boxes   2.772e-06   5.191e-08   2.772e-06   1.000000000000



========== SUMMARY (cuda, tf32) ==========
           model                  tensor     max_abs    mean_abs     max_rel       cosine_sim
---------------------------------------------------------------------------------------------
         vit_mae   reconstruction_logits   1.847e-03   1.748e-04   7.207e-04   0.999999976441
           aimv2       last_hidden_state   2.253e+00   2.005e-03   3.364e-02   0.999992942945
          d_fine     encoder_feature_map   4.338e-03   1.201e-04   2.222e-04   0.999999945733
          d_fine                  logits   2.962e+00   1.566e-01   3.188e-01   0.998815051219
          d_fine              pred_boxes   2.934e+00   1.731e-01   9.782e-01   0.850375462594
          deimv2     encoder_feature_map   2.308e-02   3.731e-05   2.229e-03   0.999999961190
          deimv2                  logits   5.575e+00   2.111e-01   5.423e-01   0.996109926608
          deimv2              pred_boxes   1.198e+00   6.097e-02   9.998e-01   0.924251090110
 pp_doclayout_v2     encoder_feature_map   1.428e-02   1.150e-04   1.251e-04   0.999999987283
 pp_doclayout_v2                  logits   3.572e+00   1.046e-01   4.211e-01   0.998505114158
 pp_doclayout_v2              pred_boxes   9.768e-01   5.563e-02   9.768e-01   0.957174794240
 pp_doclayout_v3     encoder_feature_map   3.508e-02   1.312e-04   1.379e-04   0.999999991794
 pp_doclayout_v3                  logits   4.164e+00   1.007e-01   4.216e-01   0.998889718820
 pp_doclayout_v3              pred_boxes   1.000e+00   7.107e-02   1.000e+00   0.940767513677
         rt_detr     encoder_feature_map   4.478e-03   9.512e-05   1.112e-04   0.999999985097
         rt_detr                  logits   3.300e+00   1.623e-01   3.991e-01   0.997621225686
         rt_detr              pred_boxes   9.813e-01   8.343e-02   9.813e-01   0.926827793209
      rt_detr_v2     encoder_feature_map   3.059e-02   1.684e-04   4.938e-04   0.999999973897
      rt_detr_v2                  logits   4.302e+00   1.098e-01   4.533e-01   0.998335682046
      rt_detr_v2              pred_boxes   8.968e-01   3.850e-02   8.968e-01   0.980026186848
```

On MPS I checked two approaches:
- create fp64 pos embeddings and then downcast to fp32
- create fp32 pos embeddings directly

I benchmarked against cpu fp64 embeddings. Benchmarks ran on an M5.

The differences between cpu and mps much outweigh the differences between fp64 and fp32. Fp32 errors seem even slightly smaller than fp64. So my recommendation for mps is to initialise directly on fp32 and avoid the expensive cpu->mps copy. 

Note: There also seems to be an issue with the `pp_doclayout_v2` decoder on mps as the logits have quite high errors. 



```
========== SUMMARY (mps, reference = cpu/fp64) ==========
           model                       run                 tensor     max_abs    mean_abs     max_rel       cosine_sim
----------------------------------------------------------------------------------------------------------------------
           aimv2 mps / fp64 (built on cpu)      last_hidden_state   1.616e-01   2.512e-05   2.414e-03   0.999999991243
           aimv2                mps / fp32      last_hidden_state   9.905e-02   2.433e-05   1.479e-03   0.999999994400
          d_fine mps / fp64 (built on cpu)    encoder_feature_map   1.383e-04   3.046e-06   7.084e-06   0.999999999969
          d_fine                mps / fp32    encoder_feature_map   1.323e-04   3.049e-06   6.779e-06   0.999999999969
          d_fine mps / fp64 (built on cpu)                 logits   8.082e-01   2.105e-03   8.914e-02   0.999990933781
          d_fine                mps / fp32                 logits   8.082e-01   2.099e-03   8.914e-02   0.999990934261
          d_fine mps / fp64 (built on cpu)             pred_boxes   2.755e-01   1.059e-03   9.185e-02   0.999828812957
          d_fine                mps / fp32             pred_boxes   2.755e-01   1.059e-03   9.185e-02   0.999828813294
          deimv2 mps / fp64 (built on cpu)    encoder_feature_map   2.422e-04   7.666e-07   2.337e-05   0.999999999985
          deimv2                mps / fp32    encoder_feature_map   2.422e-04   7.670e-07   2.337e-05   0.999999999985
          deimv2 mps / fp64 (built on cpu)                 logits   5.156e-03   1.512e-04   4.946e-04   0.999999999099
          deimv2                mps / fp32                 logits   5.217e-03   1.493e-04   5.004e-04   0.999999999111
          deimv2 mps / fp64 (built on cpu)             pred_boxes   2.682e-04   4.350e-06   2.237e-04   0.999999999567
          deimv2                mps / fp32             pred_boxes   2.948e-04   4.370e-06   2.459e-04   0.999999999514
 pp_doclayout_v2 mps / fp64 (built on cpu)    encoder_feature_map   9.918e-05   1.182e-06   8.690e-07   0.999999999998
 pp_doclayout_v2                mps / fp32    encoder_feature_map   9.155e-05   1.180e-06   8.021e-07   0.999999999998
 pp_doclayout_v2 mps / fp64 (built on cpu)                 logits   6.670e+00   2.701e-02   7.932e-01   0.998822744344
 pp_doclayout_v2                mps / fp32                 logits   6.670e+00   2.701e-02   7.932e-01   0.998822743898
 pp_doclayout_v2 mps / fp64 (built on cpu)             pred_boxes   9.544e-01   7.419e-03   9.544e-01   0.993605069607
 pp_doclayout_v2                mps / fp32             pred_boxes   9.544e-01   7.419e-03   9.544e-01   0.993605069102
 pp_doclayout_v3 mps / fp64 (built on cpu)    encoder_feature_map   2.594e-04   1.457e-06   1.019e-06   0.999999999999
 pp_doclayout_v3                mps / fp32    encoder_feature_map   2.594e-04   1.455e-06   1.019e-06   0.999999999999
 pp_doclayout_v3 mps / fp64 (built on cpu)                 logits   1.454e-04   6.536e-06   1.453e-05   0.999999999999
 pp_doclayout_v3                mps / fp32                 logits   1.469e-04   6.080e-06   1.467e-05   0.999999999999
 pp_doclayout_v3 mps / fp64 (built on cpu)             pred_boxes   1.603e-05   5.019e-07   1.603e-05   0.999999999998
 pp_doclayout_v3                mps / fp32             pred_boxes   1.264e-05   4.571e-07   1.264e-05   0.999999999999
         rt_detr mps / fp64 (built on cpu)    encoder_feature_map   1.121e-04   5.403e-06   2.789e-06   0.999999999965
         rt_detr                mps / fp32    encoder_feature_map   1.106e-04   5.406e-06   2.753e-06   0.999999999965
         rt_detr mps / fp64 (built on cpu)                 logits   1.727e-03   4.016e-05   2.088e-04   0.999999999846
         rt_detr                mps / fp32                 logits   1.737e-03   4.028e-05   2.100e-04   0.999999999847
         rt_detr mps / fp64 (built on cpu)             pred_boxes   6.214e-04   3.665e-06   6.214e-04   0.999999999184
         rt_detr                mps / fp32             pred_boxes   6.248e-04   3.650e-06   6.248e-04   0.999999999182
      rt_detr_v2 mps / fp64 (built on cpu)    encoder_feature_map   1.211e-04   1.265e-06   1.957e-06   0.999999999999
      rt_detr_v2                mps / fp32    encoder_feature_map   1.230e-04   1.267e-06   1.987e-06   0.999999999999
      rt_detr_v2 mps / fp64 (built on cpu)                 logits   5.379e-04   1.034e-05   5.659e-05   0.999999999995
      rt_detr_v2                mps / fp32                 logits   5.150e-04   1.002e-05   5.418e-05   0.999999999995
      rt_detr_v2 mps / fp64 (built on cpu)             pred_boxes   1.439e-05   2.040e-07   1.439e-05   0.999999999999
      rt_detr_v2                mps / fp32             pred_boxes   1.296e-05   1.954e-07   1.296e-05   0.999999999999
         vit_mae mps / fp64 (built on cpu)  reconstruction_logits   5.126e-06   8.356e-07   1.986e-06   1.000000000000
         vit_mae                mps / fp32  reconstruction_logits   5.007e-06   8.441e-07   1.939e-06   1.000000000000
```