src/autograd/reduce_fold_LAWS.bend

bendtorch-core@0.1.0.0 · 94 lines · 7 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/replayed.bend as Replayed
7import bendtorch-core@0.1.0.0/src/tensor/contracts/fold.bend as Fold
8import bendtorch-core@0.1.0.0/src/tensor/contracts/LAWS.bend as Contracts
9import bendtorch-core@0.1.0.0/src/autograd/reduce.bend as Reduce
10import bendtorch-core@0.1.0.0/src/autograd/reduce_LAWS.bend as ReduceLaws
11import bendtorch-core@0.1.0.0/src/autograd/reduce_fold.bend as ReduceFold
12
13# What a reduce cell promises, stated only through the result: it gives
14# `expected`, and the returned contribution still holds every word.
15type FiberContract<-input: Shape.Shape, -expected: Maybe<&2, F32>, -word: U32 -> F32,
16  -result: Core.Dynamic & Result<&2, &1, Core.Error, F32>> is Type:
17  FiberContract{
18    value: {Fold.value(result) == expected : Maybe<&2, F32>},
19    kept: @+other: U32 -> @inside: {Nat.is_lt(U32.to_nat(other), Shape.logical_numel(Shape.dimensions(input))) == True{} : Bool} ->
20      {Fold.input_point(result, other) == Some{word(other)} : Maybe<&2, F32>}}
21
22# Ordered-fold family, reduce instance: for a plan whose output has the
23# contribution's shape, every target index and every replayed contribution,
24# the reduce cell gives `ReduceLaws.fiber` over the contribution's words in
25# index order, and returns the contribution with every word kept.
26law reduce_fold_fiber:
27  for +plan: Shape.BroadcastPlan
28  for same: {Shape.extents(Shape.output_shape(plan)) == Shape.extents(Shape.right_shape(plan)) : List<&2, Nat>}
29  for +index: U32
30  for +fill: F32
31  for +ops: List<&2, StorageLaws.Op>
32  for -fits: {Lifecycle.logical_length(Replayed.storage(Shape.right_shape(plan), fill, ops)) ==
33    Shape.logical_numel(Shape.dimensions(Shape.right_shape(plan))) : Nat}
34  FiberContract<Shape.right_shape(plan),
35    ReduceLaws.fiber(plan,
36      Contracts.words(Shape.right_shape(plan), fill, ops, Shape.logical_numel(Shape.dimensions(Shape.right_shape(plan)))),
37      index, Shape.logical_numel(Shape.dimensions(Shape.right_shape(plan)))),
38    other => Replayed.value(Shape.right_shape(plan), fill, ops, other),
39    ReduceFold.cell(plan, index,
40      Core.Packed{Shape.right_shape(plan), Core.TTensor{Replayed.storage(Shape.right_shape(plan), fill, ops), fits}})>
41
42def output_dimensions(result: Core.Dynamic & Result<&2, &1, Reduce.Error, Core.Dynamic>) -> Maybe<&2, List<&2, U32>>:
43  (input, outcome) = result
44  match outcome:
45    case Fail{error}:
46      None{}
47    case Done{Core.Packed{shape, tensor}}:
48      Some{Shape.dimensions(shape)}
49
50def output_point(result: Core.Dynamic & Result<&2, &1, Reduce.Error, Core.Dynamic>, +index: U32) -> Maybe<&2, F32>:
51  (input, outcome) = result
52  match outcome:
53    case Fail{error}:
54      None{}
55    case Done{output}:
56      Core.dynamic_read_value(Core.read_dynamic(output, index))
57
58def input_point(result: Core.Dynamic & Result<&2, &1, Reduce.Error, Core.Dynamic>, +index: U32) -> Maybe<&2, F32>:
59  (input, outcome) = result
60  Core.dynamic_read_value(Core.read_dynamic(input, index))
61
62# What reduce_to promises, stated only through the result: it succeeds with
63# the target's dimensions, every in-bounds target cell holds `expected`, and
64# the returned contribution still holds every word.
65type ReduceToContract<-target: Shape.Shape, -input: Shape.Shape, -expected: U32 -> Maybe<&2, F32>, -word: U32 -> F32,
66  -result: Core.Dynamic & Result<&2, &1, Reduce.Error, Core.Dynamic>> is Type:
67  ReduceToContract{
68    successful: {ReduceLaws.succeeded(result) == True{} : Bool},
69    dimensions: {output_dimensions(result) == Some{Shape.dimensions(target)} : Maybe<&2, List<&2, U32>>},
70    value: @+index: U32 -> @inside: {Nat.is_lt(U32.to_nat(index), Shape.logical_numel(Shape.dimensions(target))) == True{} : Bool} ->
71      {output_point(result, index) == expected(index) : Maybe<&2, F32>},
72    kept: @+other: U32 -> @inside: {Nat.is_lt(U32.to_nat(other), Shape.logical_numel(Shape.dimensions(input))) == True{} : Bool} ->
73      {input_point(result, other) == Some{word(other)} : Maybe<&2, F32>}}
74
75# reduce_to through the ordered-fold family: when the target broadcasts to
76# exactly the contribution's dimensions, reducing a replayed contribution
77# succeeds with the target's dimensions, every target cell holds
78# `ReduceLaws.fiber` over the contribution's words in index order, the value
79# `reduce_fold_fiber` states for `ReduceFold.cell`, and the returned
80# contribution keeps every word.
81law reduce_to_fiber:
82  for +target: Shape.Shape
83  for +shape: Shape.Shape
84  for +plan: Shape.BroadcastPlan
85  for planned: {Shape.broadcast_checked(target, shape) == Done{plan} : Result<&2, &2, Shape.Error, Shape.BroadcastPlan>}
86  for same: {Shape.dimensions(Shape.output_shape(plan)) == Shape.dimensions(shape) : List<&2, U32>}
87  for +fill: F32
88  for +ops: List<&2, StorageLaws.Op>
89  for -fits: {Lifecycle.logical_length(Replayed.storage(shape, fill, ops)) == Shape.logical_numel(Shape.dimensions(shape)) : Nat}
90  ReduceToContract<target, shape,
91    index => ReduceLaws.fiber(plan, Contracts.words(shape, fill, ops, Shape.logical_numel(Shape.dimensions(shape))),
92      index, Shape.logical_numel(Shape.dimensions(shape))),
93    other => Replayed.value(shape, fill, ops, other),
94    Reduce.reduce_to(Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, ops), fits}}, target)>