def Blk.mask_cols

bend-tensors@0.0.0.2 · bend_tensors.bend:399

import Base

399# columns at index >= k set to 0.0
400def Blk.mask_cols(k: Nat, m: Blk) -> Blk:
401  match k m:
402    case 0n Blk{a00, a01, a02, a03, a10, a11, a12, a13, a20, a21, a22, a23, a30, a31, a32, a33}:
403      Blk{0.0, 0.0, 0.0, 0.0,
404      0.0, 0.0, 0.0, 0.0,
405      0.0, 0.0, 0.0, 0.0,
406      0.0, 0.0, 0.0, 0.0}
407    case 1n Blk{a00, a01, a02, a03, a10, a11, a12, a13, a20, a21, a22, a23, a30, a31, a32, a33}:
408      Blk{a00, 0.0, 0.0, 0.0,
409      a10, 0.0, 0.0, 0.0,
410      a20, 0.0, 0.0, 0.0,
411      a30, 0.0, 0.0, 0.0}
412    case 2n Blk{a00, a01, a02, a03, a10, a11, a12, a13, a20, a21, a22, a23, a30, a31, a32, a33}:
413      Blk{a00, a01, 0.0, 0.0,
414      a10, a11, 0.0, 0.0,
415      a20, a21, 0.0, 0.0,
416      a30, a31, 0.0, 0.0}
417    case 3n Blk{a00, a01, a02, a03, a10, a11, a12, a13, a20, a21, a22, a23, a30, a31, a32, a33}:
418      Blk{a00, a01, a02, 0.0,
419      a10, a11, a12, 0.0,
420      a20, a21, a22, 0.0,
421      a30, a31, a32, 0.0}
422    case _ m2:
423      m2