src/tensor/contracts/broadcast.bend

bendtorch-core@0.1.0.0 · 43 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/storage/LAWS.bend as StorageLaws
4import bendtorch-core@0.1.0.0/src/tensor/core.bend as Core
5import bendtorch-core@0.1.0.0/src/tensor/elementwise.bend as Elementwise
6import bendtorch-core@0.1.0.0/src/tensor/contracts/replayed.bend as Replayed
7import bendtorch-core@0.1.0.0/src/tensor/contracts/pointwise.bend as Pointwise
8
9def left_walk(-left: List<&2, Nat>, -right: List<&2, Nat>, -output: List<&2, Nat>,
10  map: Shape.Index<left, right, output>, +index: Nat) -> Nat:
11  match map:
12    case Shape.Indexed{rows, exact}:
13      Shape.walk_left(rows, index)
14    case Shape.Unindexed{zero}:
15      0n
16
17def right_walk(-left: List<&2, Nat>, -right: List<&2, Nat>, -output: List<&2, Nat>,
18  map: Shape.Index<left, right, output>, +index: Nat) -> Nat:
19  match map:
20    case Shape.Indexed{rows, exact}:
21      Shape.walk_right(rows, index)
22    case Shape.Unindexed{zero}:
23      0n
24
25def left_map(plan: Shape.BroadcastPlan, index: Nat) -> Nat:
26  +plan = plan
27  left_walk(Shape.extents(Shape.left_shape(plan)), Shape.extents(Shape.right_shape(plan)), Shape.extents(Shape.output_shape(plan)),
28    Shape.index_map(plan), index)
29
30def right_map(plan: Shape.BroadcastPlan, index: Nat) -> Nat:
31  +plan = plan
32  right_walk(Shape.extents(Shape.left_shape(plan)), Shape.extents(Shape.right_shape(plan)), Shape.extents(Shape.output_shape(plan)),
33    Shape.index_map(plan), index)
34
35def binary(+operation: Elementwise.Operation, +plan: Shape.BroadcastPlan, left: Core.Dynamic, right: Core.Dynamic) ->
36  (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, Core.Dynamic>:
37  Pointwise.run(~Elementwise.Operation, ~Elementwise.scalar, ~Shape.BroadcastPlan, ~left_map, ~right_map,
38    operation, plan, Shape.output_shape(plan), left, right)
39
40def cell(+operation: Elementwise.Operation, +plan: Shape.BroadcastPlan, +left_fill: F32, +left_ops: List<&2, StorageLaws.Op>,
41  +right_fill: F32, +right_ops: List<&2, StorageLaws.Op>, +index: Nat) -> F32:
42  Elementwise.scalar(operation, Replayed.value(Shape.left_shape(plan), left_fill, left_ops, U32.from_nat(left_map(plan, index))),
43    Replayed.value(Shape.right_shape(plan), right_fill, right_ops, U32.from_nat(right_map(plan, index))))