1	import Base
2	import ../shape/core.bend as Shape
3	import ../tensor/core.bend as Core
4	import ../tensor/elementwise.bend as Elementwise
5	import ../tensor/elementwise_totality.bend as Totality
6	import ../tensor/projection_LAWS.bend as Projection
7	import ./broadcast.bend as Broadcast
8	
9	def owner(pair: Core.Dynamic & Result<&2, &1, Broadcast.Error, Core.Dynamic>) -> Core.Dynamic:
10	  (input, result) = pair
11	  input
12	
13	def succeeded(pair: Core.Dynamic & Result<&2, &1, Broadcast.Error, Core.Dynamic>) -> Bool:
14	  (input, result) = pair
15	  Result.is_done(&2, &1, Broadcast.Error, Core.Dynamic, result)
16	
17	# The input value that the specification assigns to one output position: the
18	# input offset of its right-aligned broadcast projection.
19	def source_value(input: Core.Dynamic, offsets: Maybe<&1, U32 & U32>) -> Maybe<&2, F32>:
20	  match offsets:
21	    case None{}:
22	      None{}
23	    case Some{pair}:
24	      (input_offset, target_offset) = pair
25	      Core.checked_observe_dynamic(input, input_offset)
26	
27	def planned_source(input: Core.Dynamic, +index: U32, plan: Result<&2, &2, Shape.Error, Shape.BroadcastPlan>) -> Maybe<&2, F32>:
28	  match plan:
29	    case Fail{error}:
30	      None{}
31	    case Done{+plan}:
32	      source_value(input, Projection.offset_spec(plan, U32.to_nat(index)))
33	
34	def MappingWith(-input: Core.Dynamic, +target: Shape.Shape, plan: Result<&2, &2, Shape.Error, Shape.BroadcastPlan>,
35	  pair: Core.Dynamic & Result<&2, &1, Broadcast.Error, Core.Dynamic>) -> Type:
36	  (retained, result) = pair
37	  match result:
38	    case Fail{error}:
39	      Unit
40	    case Done{output}:
41	      {Totality.dimensions(output) == Shape.dimensions(target) : List<&2, U32>} &
42	        (@+index: U32 -> @inside: Shape.LT(U32.to_nat(index), Shape.logical_numel(Shape.dimensions(target))) ->
43	          {Core.observe_dynamic(output, index) == planned_source(input, index, plan) : Maybe<&2, F32>})
44	
45	def Mapping(input: Core.Dynamic, +target: Shape.Shape, pair: Core.Dynamic & Result<&2, &1, Broadcast.Error, Core.Dynamic>) -> Type:
46	  MappingWith(input, target, Shape.broadcast_checked(Totality.dynamic_shape(input), target), pair)
47	
48	type Retained<-input: Core.Dynamic, -pair: Core.Dynamic & Result<&2, &1, Broadcast.Error, Core.Dynamic>> is Data:
49	  Retained{-kept: Totality.SameSource(input, owner(pair))}
50	
51	type MappingModel<-input: Core.Dynamic, -target: Shape.Shape, -pair: Core.Dynamic & Result<&2, &1, Broadcast.Error, Core.Dynamic>> is Data:
52	  MappingModel{-mapping: Mapping(input, target, pair)}
53	
54	type Success<-pair: Core.Dynamic & Result<&2, &1, Broadcast.Error, Core.Dynamic>> is Data:
55	  Success{-successful: {succeeded(pair) == True{} : Bool}}
56	
57	law keep_left_scalar:
58	  for +left: F32
59	  for +right: F32
60	  {Elementwise.scalar(Elementwise.KeepLeft{}, left, right) == left : F32}
61	
62	# The input owner comes back with the same dimensions and observations.
63	law broadcast_retains_input:
64	  for input: Core.Dynamic
65	  for +target: Shape.Shape
66	  Retained<input, Broadcast.broadcast_to(input, target)>
67	
68	# On success the output has the target's dimensions, and each in-bounds output
69	# position observes the input at the offset that the independent broadcast
70	# projection specification assigns to it.
71	law broadcast_mapping:
72	  for input: Core.Dynamic
73	  for +target: Shape.Shape
74	  MappingModel<input, target, Broadcast.broadcast_to(input, target)>
75	
76	# A request whose right-aligned broadcast with the input shape is the target
77	# itself always succeeds.
78	law broadcast_valid_success:
79	  for input: Core.Dynamic
80	  for +target: Shape.Shape
81	  for +alignment: Shape.Broadcast<Shape.extents(Totality.dynamic_shape(input)), Shape.extents(target), Shape.extents(target)>
82	  Success<Broadcast.broadcast_to(input, target)>
83	
84	# When the broadcast planner rejects the input and target shapes, broadcast_to
85	# fails with exactly that planner error and returns the input owner. For example [2] to
86	# [0] fails with IncompatibleBroadcast instead of returning an empty tensor.
87	law broadcast_shape_error:
88	  for input: Core.Dynamic
89	  for +target: Shape.Shape
90	  for error: Shape.Error
91	  for rejected: {Shape.broadcast_checked(Totality.dynamic_shape(input), target) == Fail{error} : Result<&2, &2, Shape.Error, Shape.BroadcastPlan>}
92	  {Broadcast.broadcast_to(input, target) == (input, Fail{Broadcast.TensorError{Core.ShapeError{error}}}) : Core.Dynamic & Result<&2, &1, Broadcast.Error, Core.Dynamic>}
93	
94	# When the planned broadcast is not the target itself (for example [3] to
95	# [1]), broadcast_to fails with NotBroadcastable and returns the input owner.
96	law broadcast_not_target:
97	  for input: Core.Dynamic
98	  for +target: Shape.Shape
99	  for +plan: Shape.BroadcastPlan
100	  for planned: {Shape.broadcast_checked(Totality.dynamic_shape(input), target) == Done{plan} : Result<&2, &2, Shape.Error, Shape.BroadcastPlan>}
101	  for differ: {Shape.dimensions(Shape.output_shape(plan)) != Shape.dimensions(target) : List<&2, U32>}
102	  {Broadcast.broadcast_to(input, target) == (input, Fail{Broadcast.NotBroadcastable{}}) : Core.Dynamic & Result<&2, &1, Broadcast.Error, Core.Dynamic>}