~/jendhub

src/autograd/broadcast_pointwise_LAWS.bend

bendtorch-core@0.1.0.0 · 112 lines · 9 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/lifecycle.bend as Lifecycle
6import bendtorch-core@0.1.0.0/src/tensor/contracts/LAWS.bend as Contracts
7import bendtorch-core@0.1.0.0/src/tensor/contracts/replayed.bend as Replayed
8import bendtorch-core@0.1.0.0/src/tensor/contracts/broadcast.bend as PointwiseBroadcast
9import bendtorch-core@0.1.0.0/src/autograd/broadcast.bend as Broadcast
10
11def succeeded(result: Core.Dynamic & Result<&2, &1, Broadcast.Error, Core.Dynamic>) -> Bool:
12  (input, outcome) = result
13  Result.is_done(&2, &1, Broadcast.Error, Core.Dynamic, outcome)
14
15def output_dimensions(result: Core.Dynamic & Result<&2, &1, Broadcast.Error, Core.Dynamic>) -> Maybe<&2, List<&2, U32>>:
16  (input, outcome) = result
17  match outcome:
18    case Fail{error}:
19      None{}
20    case Done{Core.Packed{shape, tensor}}:
21      Some{Shape.dimensions(shape)}
22
23def output_point(result: Core.Dynamic & Result<&2, &1, Broadcast.Error, Core.Dynamic>, +index: U32) -> Maybe<&2, F32>:
24  (input, outcome) = result
25  match outcome:
26    case Fail{error}:
27      None{}
28    case Done{output}:
29      Core.dynamic_read_value(Core.read_dynamic(output, index))
30
31def input_point(result: Core.Dynamic & Result<&2, &1, Broadcast.Error, Core.Dynamic>, +index: U32) -> Maybe<&2, F32>:
32  (input, outcome) = result
33  Core.dynamic_read_value(Core.read_dynamic(input, index))
34
35# What broadcast_to promises, stated only through the result: it succeeds
36# with the target's dimensions, every in-bounds output index i holds cell(i),
37# and the returned input still holds every word.
38type BroadcastToContract<-target: Shape.Shape, -input: Shape.Shape, -cell: Nat -> F32, -word: U32 -> F32,
39  -result: Core.Dynamic & Result<&2, &1, Broadcast.Error, Core.Dynamic>> is Type:
40  BroadcastToContract{
41    successful: {succeeded(result) == True{} : Bool},
42    dimensions: {output_dimensions(result) == Some{Shape.dimensions(target)} : Maybe<&2, List<&2, U32>>},
43    value: @+index: U32 -> @inside: {Nat.is_lt(U32.to_nat(index), Shape.logical_numel(Shape.dimensions(target))) == True{} : Bool} ->
44      {output_point(result, index) == Some{cell(U32.to_nat(index))} : Maybe<&2, F32>},
45    kept: @+other: U32 -> @inside: {Nat.is_lt(U32.to_nat(other), Shape.logical_numel(Shape.dimensions(input))) == True{} : Bool} ->
46      {input_point(result, other) == Some{word(other)} : Maybe<&2, F32>}}
47
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)))])
57
58# What the transport form of broadcast_to_replayed asks of P: P holds for the
59# input, replayed with any history that keeps its words, next to an output of
60# the plan's shape, replayed over +0.0 with any history whose words are
61# `cells`.
62def BroadcastWith(+plan: Shape.BroadcastPlan, +shape: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>,
63  -P: Core.Dynamic & Result<&2, &1, Broadcast.Error, Core.Dynamic> -> Type) -> Type:
64  @+kept: List<&2, StorageLaws.Op> ->
65  @kept_words: {Contracts.words(shape, fill, kept, Shape.logical_numel(Shape.dimensions(shape))) ==
66    Contracts.words(shape, fill, ops, Shape.logical_numel(Shape.dimensions(shape))) : List<&2, F32>} ->
67  @-input_fits: {Lifecycle.logical_length(Replayed.storage(shape, fill, kept)) == Shape.logical_numel(Shape.dimensions(shape)) : Nat} ->
68  @+output_ops: List<&2, StorageLaws.Op> ->
69  @values: {Contracts.words(Shape.output_shape(plan), 0.0, output_ops, Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan)))) ==
70    cells(plan, shape, fill, ops, Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan)))) : List<&2, F32>} ->
71  @-output_fits: {Lifecycle.logical_length(Replayed.storage(Shape.output_shape(plan), 0.0, output_ops)) ==
72    Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan))) : Nat} ->
73    P((Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, kept), input_fits}},
74      Done{Core.Packed{Shape.output_shape(plan), Core.TTensor{Replayed.storage(Shape.output_shape(plan), 0.0, output_ops), output_fits}}}))
75
76# Pointwise family, broadcast_to instance: when the input shape broadcasts to
77# exactly the target, broadcasting a replayed input succeeds with the target's
78# dimensions, output index i holds the input word at left_map(plan, i), its
79# right-aligned broadcast projection, and the returned input keeps every word.
80law broadcast_to_pointwise:
81  for +shape: Shape.Shape
82  for +target: Shape.Shape
83  for +plan: Shape.BroadcastPlan
84  for planned: {Shape.broadcast_checked(shape, target) == Done{plan} : Result<&2, &2, Shape.Error, Shape.BroadcastPlan>}
85  for same: {Shape.dimensions(Shape.output_shape(plan)) == Shape.dimensions(target) : List<&2, U32>}
86  for +fill: F32
87  for +ops: List<&2, StorageLaws.Op>
88  for -fits: {Lifecycle.logical_length(Replayed.storage(shape, fill, ops)) == Shape.logical_numel(Shape.dimensions(shape)) : Nat}
89  BroadcastToContract<target, shape,
90    index => Replayed.value(shape, fill, ops, U32.from_nat(PointwiseBroadcast.left_map(plan, index))),
91    other => Replayed.value(shape, fill, ops, other),
92    Broadcast.broadcast_to(Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, ops), fits}}, target)>
93
94# The broadcast_to result is replayed: for every P, if P holds for every
95# result that hands back the input with its words and gives an output of the
96# plan's shape whose words are `cells`, then P holds for broadcast_to. The
97# output word at k is the input word at left_map(plan, k), whatever split of
98# the output into parallel halves the build takes, and the output is a
99# replayed tensor, so a later law about a replayed contribution, such as
100# reduce_to_fiber, applies to it.
101law broadcast_to_replayed:
102  for +shape: Shape.Shape
103  for +target: Shape.Shape
104  for +plan: Shape.BroadcastPlan
105  for planned: {Shape.broadcast_checked(shape, target) == Done{plan} : Result<&2, &2, Shape.Error, Shape.BroadcastPlan>}
106  for same: {Shape.dimensions(Shape.output_shape(plan)) == Shape.dimensions(target) : List<&2, U32>}
107  for +fill: F32
108  for +ops: List<&2, StorageLaws.Op>
109  for -fits: {Lifecycle.logical_length(Replayed.storage(shape, fill, ops)) == Shape.logical_numel(Shape.dimensions(shape)) : Nat}
110  for -P: Core.Dynamic & Result<&2, &1, Broadcast.Error, Core.Dynamic> -> Type
111  for h: BroadcastWith(plan, shape, fill, ops, P)
112  P(Broadcast.broadcast_to(Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, ops), fits}}, target))