def MatT.mask

bend-tensors@0.0.0.2 · bend_tensors.bend:1267

import Base

1267# rows at index >= rows and columns at index >= cols set to 0.0
1268def MatT.mask(+r: Nat, +c: Nat, m: MatT(r, c), +i0: Nat, +j0: Nat, +rows: Nat, +cols: Nat) -> MatT(r, c):
1269  match r c:
1270    case 0n 0n:
1271      MatB{b} = m
1272      MatB{Blk.mask(Nat.sub(rows, i0), Nat.sub(cols, j0), b)}
1273    case 1n+p 0n:
1274      MatR{t, b} = m
1275      x y = MatT.mask(p, 0n, t, i0, j0, rows, cols) MatT.mask(p, 0n, b, Nat.add(i0, VecT.size(p)), j0, rows, cols)
1276      MatR{x, y}
1277    case 0n 1n+q:
1278      MatC{l, rt} = m
1279      x y = MatT.mask(0n, q, l, i0, j0, rows, cols) MatT.mask(0n, q, rt, i0, Nat.add(j0, VecT.size(q)), rows, cols)
1280      MatC{x, y}
1281    case 1n+p 1n+q:
1282      MatQ{a, b, c2, d} = m
1283      +i1 = Nat.add(i0, VecT.size(p))
1284      +j1 = Nat.add(j0, VecT.size(q))
1285      w x y z = MatT.mask(p, q, a, i0, j0, rows, cols) MatT.mask(p, q, b, i0, j1, rows, cols) MatT.mask(p, q, c2, i1, j0, rows, cols) MatT.mask(p, q, d, i1, j1, rows, cols)
1286      MatQ{w, x, y, z}