Apply Qwen vision RoPE per row across table boundaries

tools/issue108-reference.py

@@ -0,0 +1,55 @@
+"""Pinned attention chunk RoPE oracle, with independent per-row controls."""
+import hashlib
+import json
+from pathlib import Path
+import subprocess
+import sys
+
+checkout, output = map(Path, sys.argv[1:3])
+assert subprocess.check_output(['git', '-C', str(checkout), 'rev-parse', 'HEAD'], text=True).strip() == '9882703f3105363ddc37eca9f97aa09a1d387112'
+sys.path.insert(0, str(checkout))
+import mlx.core as mx
+import numpy as np
+from mtplx.models.qwen4_exp import _vision_chunk_cos_sin, _apply_partial_rope, _build_mrope_axes, _rope_cos_sin, _rope_inv_freq_and_scaling, TextArgs
+
+output.mkdir(parents=True, exist_ok=True)
+manifest = []
+for dim, rotary, sections, heads in [(8, 8, [2,1,1], 2), (256,64,[11,11,10],24)]:
+    axes = mx.array(_build_mrope_axes(sections, True), dtype=mx.int32)
+    freq, _ = _rope_inv_freq_and_scaling(TextArgs(head_dim=dim, rope_parameters=dict(rope_type="default", rope_theta=10000000, partial_rotary_factor=rotary/dim)))
+    short = [[0,1,2,2],[0,1,3,3],[0,1,4,4]]
+    full = [r + [3,4,5,6] for r in short]
+    equal = [list(range(4))] * 3
+    for mode, table, start, rows, delta in [
+        ('inside',short,0,4,-1), ('outside',short,4,2,-1),
+        ('straddle',short,2,4,-1), ('full_text_tail',full,6,4,-1),
+        ('equal',equal,2,4,0), ('none',None,4,2,-1), ('empty',[[],[],[]],4,2,-1),
+    ]:
+        for scaling in [1.,1.25]:
+            t = None if table is None else mx.array(table,dtype=mx.int32)
+            cos, sin = _vision_chunk_cos_sin(t,delta,start,rows,freq,axes,scaling)
+            # Independent row-position construction; no chunk helper is used.
+            pos3 = np.stack([np.array(table)[:,p] if table is not None and p<len(table[0]) else np.full(3,p+delta) for p in range(start,start+rows)],axis=1)
+            angles = pos3[np.asarray(axes)].T.astype(np.float32) * np.asarray(freq)[None,:]
+            angles = np.concatenate([angles,angles],axis=1)
+            np.testing.assert_allclose(np.asarray(cos),np.cos(angles)*scaling,rtol=3e-6,atol=3e-6)
+            np.testing.assert_allclose(np.asarray(sin),np.sin(angles)*scaling,rtol=3e-6,atol=3e-6)
+            if mode in ('full_text_tail','equal'):
+                plain = _rope_cos_sin(mx.arange(start+delta,start+rows+delta,dtype=mx.int32),freq,scaling)
+                assert all(mx.array_equal(a,b).item() for a,b in zip((cos,sin),plain))
+            for dtype_name,dtype in [('f32',mx.float32),('bf16',mx.bfloat16)]:
+                key = f'{dim}-{mode}-{scaling}-{dtype_name}'
+                record = dict(key=key,dim=dim,rotary=rotary,sections=sections,heads=heads,mode=mode,table=table,start=start,rows=rows,delta=delta,scaling=scaling,dtype=dtype_name)
+                arrays = [('cos',cos),('sin',sin)]
+                for part,h,seed in [('q',heads,3),('k',1,11)]:
+                    shape = (1,rows,h,dim)
+                    raw = (((np.arange(np.prod(shape),dtype=np.int64)*17+seed)%97-48)/32).astype(np.float32).reshape(shape)
+                    value = _apply_partial_rope(mx.array(raw).astype(dtype),cos,sin)
+                    arrays.append((part,value))
+                for part,value in arrays:
+                    data=np.asarray(value.astype(mx.float32)).tobytes()
+                    (output/f'{key}-{part}.f32').write_bytes(data)
+                    record[part+'_sha256']=hashlib.sha256(data).hexdigest()
+                manifest.append(record)
+                print(json.dumps(record),flush=True)
+(output/'manifest.json').write_text(json.dumps(manifest))