src/autograd/reduce_fold.bend

bendtorch-core@0.1.0.0 · 30 lines · 6 definitions

1import Base
2import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
3import bendtorch-core@0.1.0.0/src/tensor/core.bend as Core
4import bendtorch-core@0.1.0.0/src/tensor/contracts/fold.bend as Fold
5import bendtorch-core@0.1.0.0/src/tensor/contracts/sum.bend as Sum
6import bendtorch-core@0.1.0.0/src/tensor/contracts/broadcast.bend as Broadcast
7
8# Ordered-fold family, reduce instance: target cell `index` of reduce_to sums,
9# in increasing position, the contribution words whose target offset is
10# `index`.
11
12type Cell is Data:
13  Cell{plan: Shape.BroadcastPlan, index: U32}
14
15def target_offset(plan: Shape.BroadcastPlan, position: Nat) -> Nat:
16  Broadcast.left_map(plan, position)
17
18def contribution_offset(plan: Shape.BroadcastPlan, position: Nat) -> Nat:
19  Broadcast.right_map(plan, position)
20
21def source(cell: Cell, position: Nat) -> Nat:
22  Cell{+plan, index} = cell
23  contribution_offset(plan, position)
24
25def lands(cell: Cell, position: Nat) -> Bool:
26  Cell{+plan, +index} = cell
27  Nat.is_eq(target_offset(plan, position), U32.to_nat(index))
28
29def cell(+plan: Shape.BroadcastPlan, +index: U32, input: Core.Dynamic) -> Core.Dynamic & Result<&2, &1, Core.Error, F32>:
30  Fold.run(~Unit, ~Sum.add, ~Cell, ~source, ~lands, Unit{}, Cell{plan, index}, Shape.elements(Shape.right_shape(plan)), input)