def fiber

bendtorch-core@0.1.0.0 · src/autograd/reduce.bend:64

import Base
import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape

64# The words of the fiber of `target` below these rows, added to
65# `accumulator` in increasing position. `extents` lists the contribution
66# extent of each row and `base` is the contribution offset the rows above
67# chose. The first row walks `fuel` digits from `digit`, each at offset
68# base + right * digit, and the rows below get `target % left` (all of
69# `target` on a reduced axis, as `Nat.mod` by 0 keeps it).
70def fiber(+rows: List<&2, Shape.Row>, +extents: List<&2, Nat>, fuel: Nat, +target: Nat, +base: Nat, +digit: Nat,
71  accumulator: F32, array: Array<F32>) -> F32 & Array<F32>:
72  match rows extents fuel:
73    case Nil{} _ _:
74      leaf_add(accumulator, Array.get(F32, array, U32.from_nat(base)))
75    case Shape.Row{+stride, +left, +right} <> +rest +extent <> +others 0n:
76      (accumulator, array)
77    case Shape.Row{+stride, +left, +right} <> +rest +extent <> +others 1n+remaining:
78      +inner = Nat.mod(target, left)
79      walk_next(value => kept => fiber(rows, extents, remaining, target, base, 1n+digit, value, kept),
80        fiber(rest, others, first_span(rest, others), inner, Nat.add(base, Nat.mul(right, digit)), first_digit(rest, inner),
81          accumulator, array))
82    case Shape.Row{+stride, +left, +right} <> +rest Nil{} _:
83      (accumulator, array)