def cells

bendtorch-core@0.1.0.0 · src/autograd/broadcast_pointwise_LAWS.bend:48

import Base
import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
import bendtorch-core@0.1.0.0/src/storage/LAWS.bend as StorageLaws
import bendtorch-core@0.1.0.0/src/tensor/contracts/replayed.bend as Replayed
import bendtorch-core@0.1.0.0/src/tensor/contracts/broadcast.bend as PointwiseBroadcast

48# The words of broadcast_to's output in index order: word k is the input
49# word at left_map(plan, k), its right-aligned broadcast projection.
50def cells(+plan: Shape.BroadcastPlan, +shape: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>, +count: Nat) -> List<&2, F32>:
51  match count:
52    case 0n:
53      Nil{}
54    case 1n+previous:
55      List.append(&2, F32, cells(plan, shape, fill, ops, previous),
56        [Replayed.value(shape, fill, ops, U32.from_nat(PointwiseBroadcast.left_map(plan, previous)))])