1	import Base
2	import ../shape/core.bend as Shape
3	import ../tensor/core.bend as Core
4	import ../tensor/layout.bend as Layout
5	import ./reduce.bend as Reduce
6	
7	# The PyTorch-facing reducer for a broadcast contribution: the gradient
8	# PyTorch's autograd engine gives an operand of a broadcast operation. When the
9	# target has the contribution's dimensions nothing is reduced and the output is
10	# a clone of the contribution, so every word, including -0.0 and NaN payloads,
11	# is kept. Otherwise the result is exactly Reduce.reduce_to's ordered fold.
12	
13	def cloned(+shape: Shape.Shape, pair: Core.Tensor<shape> & Core.Tensor<shape>) ->
14	  Core.Dynamic & Result<&2, &1, Reduce.Error, Core.Dynamic>:
15	  (kept, copy) = pair
16	  (Core.Packed{shape, kept}, Done{Core.Packed{shape, copy}})
17	
18	def decided(+shape: Shape.Shape, tensor: Core.Tensor<shape>, +target: Shape.Shape, same: Bool) ->
19	  Core.Dynamic & Result<&2, &1, Reduce.Error, Core.Dynamic>:
20	  match same:
21	    case False{}:
22	      Reduce.reduce_to(Core.Packed{shape, tensor}, target)
23	    case True{}:
24	      cloned(shape, Core.clone(shape, tensor))
25	
26	# The contribution owner comes back next to the result, as with reduce_to.
27	def sum_to_size(input: Core.Dynamic, +target: Shape.Shape) -> Core.Dynamic & Result<&2, &1, Reduce.Error, Core.Dynamic>:
28	  Core.Packed{+shape, tensor} = input
29	  decided(shape, tensor, target, Layout.words_same(Shape.dimensions(shape), Shape.dimensions(target)))
30	
31	def sum_to_size_dimensions_result(input: Core.Dynamic, -dimensions: List<&2, U32>, result: Shape.ConstructResult<dimensions>) ->
32	  Core.Dynamic & Result<&2, &1, Reduce.Error, Core.Dynamic>:
33	  match result:
34	    case Shape.SError{error}:
35	      (input, Fail{Reduce.TensorError{Core.ShapeError{error}}})
36	    case Shape.SBuilt{+target, preserved}:
37	      sum_to_size(input, target)
38	
39	def sum_to_size_dimensions(input: Core.Dynamic, +dimensions: List<&2, U32>) -> Core.Dynamic & Result<&2, &1, Reduce.Error, Core.Dynamic>:
40	  sum_to_size_dimensions_result(input, dimensions, Shape.construct(dimensions))