def Vec4.argmax_upto

bend-tensors@0.0.0.2 · bend_tensors.bend:361

import Base

361# the index and value of the largest of the first k elements
362def Vec4.argmax_upto(k: Nat, v: Vec4) -> U32 & F32:
363  match k v:
364    case 1n Vec4{x0, x1, x2, x3}:
365      (0, x0)
366    case 2n Vec4{x0, x1, x2, x3}:
367      Vec4.argmax.two((0, x0), (1, x1))
368    case 3n Vec4{x0, x1, x2, x3}:
369      Vec4.argmax.two(Vec4.argmax.two((0, x0), (1, x1)), (2, x2))
370    case _ v2:
371      Vec4.argmax(v2)