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:

  1. assert_numeric_metrics on values (PCC 0.9999, atol/rtol 1e-6).
  2. assert_equal on values vs torch.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.
  3. 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).
Installs
1
GitHub Stars
5
First Seen
6 days ago
Security Audits