1	import Base
2	import ../shape/core.bend as Shape
3	import ../storage/LAWS.bend as StorageLaws
4	import ../tensor/core.bend as Core
5	import ../tensor/lifecycle.bend as Lifecycle
6	import ../tensor/contracts/replayed.bend as Replayed
7	import ../tensor/contracts/fold.bend as Fold
8	import ../tensor/contracts/LAWS.bend as Contracts
9	import ./reduce.bend as Reduce
10	import ./reduce_LAWS.bend as ReduceLaws
11	import ./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.
15	type 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.
26	law 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	
42	def 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	
50	def 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	
58	def 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.
65	type 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.
81	law 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)>