def ms_blk

bendtorch-gpu@1.3.0.0 · main.bend:1181

import Base

1181# the squares of the block of 2^p cells from i summed as a balanced tree of pairs (h = 2^(p-1)), gradients written
1182def ms_blk(p: Nat, +i: U32, +h: U32, +kk: F32, y: Array<F32>, t: Array<F32>) -> MS:
1183  match p:
1184    case 0n: ms_cell(i, kk, Array.get(F32, y, i), Array.get(F32, t, i))
1185    case 1n++q: ms_then(ms_blk(q, i, U32.shr(h), kk, y, t), +s => y2 => t2 => ms_add(s, ms_blk(q, (i + h : U32), U32.shr(h), kk, y2, t2)))