def VecT.mask

bend-tensors@0.0.0.2 · bend_tensors.bend:1211

import Base

1211# elements at index >= m set to 0.0
1212def VecT.mask(+n: Nat, v: VecT(n), +i0: Nat, +m: Nat) -> VecT(n):
1213  match n:
1214    case 0n:
1215      VecL{x} = v
1216      VecL{Vec4.mask(Nat.sub(m, i0), x)}
1217    case 1n+p:
1218      VecN{l, r} = v
1219      a b = VecT.mask(p, l, i0, m) VecT.mask(p, r, Nat.add(i0, VecT.size(p)), m)
1220      VecN{a, b}