fix(rvt): AxialRotaryEmbedding preserves float16 / bfloat16 input dtype

[tritsystem]

What

RvT does not run in half precision — RuntimeError: expected m1 and m2 to have the same dtype.

AxialRotaryEmbedding.forward unpacks dtype from x but never uses it:

device, dtype, n = x.device, x.dtype, int(sqrt(x.shape[-2]))
seq = torch.linspace(-1., 1., steps = n, device = device)   # no dtype= -> float32

scales.to(x) doesn't rescue it (float32 * half -> float32), so sin / cos come back float32. In RvT those multiply the half-precision q / k inside apply_rotary_emb and the attention matmul raises.

Repro (on main)

import torch
from vit_pytorch.rvt import RvT

m = RvT(image_size = 32, patch_size = 8, num_classes = 5,
        dim = 64, depth = 2, heads = 4, mlp_dim = 128).half()

m(torch.randn(1, 3, 32, 32, dtype = torch.float16))
# RuntimeError: expected m1 and m2 to have the same dtype, but got: float != struct c10::Half

AxialRotaryEmbedding alone: float16 / bfloat16 in -> float32 sin, cos out.

Fix

Pass dtype = dtype to torch.linspace — the dtype the method already unpacks. float32 output is unchanged; float16 / bfloat16 now round-trip.

Test

tests/test_rvt.py builds RvT in float16 and bfloat16 and asserts the output shape and dtype.

• On main: 2 failed (RuntimeError: expected m1 and m2 to have the same dtype)
• With this change: 2 passed

AI disclosure

Found by a metamorphic dtype-preservation sweep that flagged AxialRotaryEmbedding's torch.linspace(..., device = device) with no dtype= and the unused dtype local, then confirmed float16 in -> float32 out. I verified the repro, the fix, and fail-on-main / pass-with-fix.