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/layout.bend as Layout
6	
7	# Materialized broadcast: an owned tensor of the requested shape whose every
8	# position holds the input value at its right-aligned broadcast projection.
9	# The elementwise parallel build writes it, KeepLeft with the input on the left
10	# and a +0.0 tensor of the target on the right, so an output above 2^15 cells
11	# is built in parallel halves.
12	
13	type Error is Data:
14	  TensorError{error: Core.Error}
15	  NotBroadcastable{}
16	
17	def finish(pair: (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, Core.Dynamic>) ->
18	  Core.Dynamic & Result<&2, &1, Error, Core.Dynamic>:
19	  ((input, zeros), result) = pair
20	  match result:
21	    case Fail{error}:
22	      (input, Fail{TensorError{error}})
23	    case Done{output}:
24	      (input, Done{output})
25	
26	def zeroed(input: Core.Dynamic, +target: Shape.Shape, +plan: Shape.BroadcastPlan,
27	  result: Result<&2, &1, Core.Error, Core.Tensor<target>>) -> Core.Dynamic & Result<&2, &1, Error, Core.Dynamic>:
28	  match result:
29	    case Fail{error}:
30	      (input, Fail{TensorError{error}})
31	    case Done{zeros}:
32	      finish(Elementwise.parallel_allocated(Elementwise.KeepLeft{}, plan, input, Core.Packed{target, zeros},
33	        Core.new(Shape.output_shape(plan), 0.0)))
34	
35	def fitted(input: Core.Dynamic, +target: Shape.Shape, +plan: Shape.BroadcastPlan, same: Bool) ->
36	  Core.Dynamic & Result<&2, &1, Error, Core.Dynamic>:
37	  match same:
38	    case False{}:
39	      (input, Fail{NotBroadcastable{}})
40	    case True{}:
41	      zeroed(input, target, plan, Core.new(target, 0.0))
42	
43	def planned(input: Core.Dynamic, +target: Shape.Shape, result: Result<&2, &2, Shape.Error, Shape.BroadcastPlan>) ->
44	  Core.Dynamic & Result<&2, &1, Error, Core.Dynamic>:
45	  match result:
46	    case Fail{error}:
47	      (input, Fail{TensorError{Core.ShapeError{error}}})
48	    case Done{+plan}:
49	      fitted(input, target, plan, Layout.words_same(Shape.dimensions(Shape.output_shape(plan)), Shape.dimensions(target)))
50	
51	# The input owner comes back unchanged next to the result. An incompatible
52	# shape pair fails with the planner's broadcast error; a compatible pair whose
53	# broadcast is larger than the target (for example [3] to [1]) fails with
54	# NotBroadcastable.
55	def broadcast_to(input: Core.Dynamic, +target: Shape.Shape) -> Core.Dynamic & Result<&2, &1, Error, Core.Dynamic>:
56	  Core.Packed{+shape, tensor} = input
57	  planned(Core.Packed{shape, tensor}, target, Shape.broadcast_checked(shape, target))
58	
59	def broadcast_to_dimensions_result(input: Core.Dynamic, -dimensions: List<&2, U32>, result: Shape.ConstructResult<dimensions>) ->
60	  Core.Dynamic & Result<&2, &1, Error, Core.Dynamic>:
61	  match result:
62	    case Shape.SError{error}:
63	      (input, Fail{TensorError{Core.ShapeError{error}}})
64	    case Shape.SBuilt{+target, preserved}:
65	      broadcast_to(input, target)
66	
67	def broadcast_to_dimensions(input: Core.Dynamic, +dimensions: List<&2, U32>) -> Core.Dynamic & Result<&2, &1, Error, Core.Dynamic>:
68	  broadcast_to_dimensions_result(input, dimensions, Shape.construct(dimensions))
69	
70	def error_code(error: Error) -> U32:
71	  match error:
72	    case TensorError{error}:
73	      Core.error_code(error)
74	    case NotBroadcastable{}:
75	      100