def Blk.mask_rows

bend-tensors@0.0.0.2 · bend_tensors.bend:373

import Base

373# rows at index >= k set to 0.0
374def Blk.mask_rows(k: Nat, m: Blk) -> Blk:
375  match k m:
376    case 0n Blk{a00, a01, a02, a03, a10, a11, a12, a13, a20, a21, a22, a23, a30, a31, a32, a33}:
377      Blk{0.0, 0.0, 0.0, 0.0,
378      0.0, 0.0, 0.0, 0.0,
379      0.0, 0.0, 0.0, 0.0,
380      0.0, 0.0, 0.0, 0.0}
381    case 1n Blk{a00, a01, a02, a03, a10, a11, a12, a13, a20, a21, a22, a23, a30, a31, a32, a33}:
382      Blk{a00, a01, a02, a03,
383      0.0, 0.0, 0.0, 0.0,
384      0.0, 0.0, 0.0, 0.0,
385      0.0, 0.0, 0.0, 0.0}
386    case 2n Blk{a00, a01, a02, a03, a10, a11, a12, a13, a20, a21, a22, a23, a30, a31, a32, a33}:
387      Blk{a00, a01, a02, a03,
388      a10, a11, a12, a13,
389      0.0, 0.0, 0.0, 0.0,
390      0.0, 0.0, 0.0, 0.0}
391    case 3n Blk{a00, a01, a02, a03, a10, a11, a12, a13, a20, a21, a22, a23, a30, a31, a32, a33}:
392      Blk{a00, a01, a02, a03,
393      a10, a11, a12, a13,
394      a20, a21, a22, a23,
395      0.0, 0.0, 0.0, 0.0}
396    case _ m2:
397      m2