index-based-ops
Installation
SKILL.md
Index-based ops (TopK, Sort, Argmax)
These ops return values + indices. emule does the math in f32, silicon in the SFPU; on value ties the index assignment can differ yet still be correct. Match the test's contract, not silicon's exact index choice.
What the test asserts (TopK; Sort/Argmax are analogous)
test_topk.py checks:
assert_numeric_metricson values (PCC 0.9999, atol/rtol 1e-6).assert_equalon values vstorch.topk(sorted=True)— exact. bf16 round-trips losslessly (CB bf16 → f32 DST → bf16), so the correct top-K values match torch exactly regardless of which tied index won. The only requirement: select the correct top-K set, sorted.- Gather cosine on indices:
cosine(torch_vals, gather(input, ttnn_idx)) > 0.99. Indices are never compared to torch directly — only via gather, which is robust to ties ⇒ indices need only be valid (point at the right values).