def finish_mapping

bendtorch-core@0.1.0.0 · src/autograd/broadcast_PROOF.bend:136

import Base
import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
import bendtorch-core@0.1.0.0/src/tensor/core.bend as Core
import bendtorch-core@0.1.0.0/src/tensor/elementwise.bend as Elementwise
import bendtorch-core@0.1.0.0/src/tensor/elementwise_totality.bend as Totality
import bendtorch-core@0.1.0.0/src/tensor/elementwise_PROOF.bend as ElementwiseProof
import bendtorch-core@0.1.0.0/src/tensor/projection_LAWS.bend as Projection
import bendtorch-core@0.1.0.0/src/autograd/broadcast.bend as Broadcast
import bendtorch-core@0.1.0.0/src/autograd/broadcast_LAWS.bend as Laws

136def finish_mapping(-input: Core.Dynamic, +target: Shape.Shape, -zeros: Core.Tensor<target>, +plan: Shape.BroadcastPlan,
137  pair: (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, Core.Dynamic>,
138  -dims: Totality.DimensionsModel<Shape.dimensions(Shape.output_shape(plan)), pair>,
139  -mapping: Elementwise.MappingModel<Elementwise.KeepLeft{}, plan, input, Core.Packed{target, zeros}, pair>,
140  -same_dims: {Shape.dimensions(Shape.output_shape(plan)) == Shape.dimensions(target) : List<&2, U32>},
141  -right_dims: {Shape.dimensions(Shape.right_shape(plan)) == Shape.dimensions(target) : List<&2, U32>}) ->
142  WithModel<input, target, Done{plan}, Broadcast.finish(pair)>:
143  ((left, right), result) = pair
144  match result:
145    case Fail{error}:
146      WithModel{Unit{}}
147    case Done{output}:
148      WithModel{(
149        Equal.trans(List<&2, U32>, Totality.dimensions(output), Shape.dimensions(Shape.output_shape(plan)), Shape.dimensions(target),
150          {\{Totality.DimensionsModel: exact => exact} : Totality.DimensionsModel<Shape.dimensions(Shape.output_shape(plan)), ((left, right), Done{output})> ->
151            {Totality.dimensions(output) == Shape.dimensions(Shape.output_shape(plan)) : List<&2, U32>}}(dims),
152          same_dims),
153        index => inside => Equal.trans(Maybe<&2, F32>, Core.observe_dynamic(output, index),
154          Elementwise.expected(Elementwise.KeepLeft{}, plan, input, Core.Packed{target, zeros}, index),
155          Laws.source_value(input, Projection.offset_spec(plan, U32.to_nat(index))),
156          {\{Elementwise.MappingModel: prefix => prefix} :
157            Elementwise.MappingModel<Elementwise.KeepLeft{}, plan, input, Core.Packed{target, zeros}, ((left, right), Done{output})> ->
158            Elementwise.PrefixMapping(Elementwise.KeepLeft{}, plan, input, Core.Packed{target, zeros},
159              Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan))), output)}(mapping)(index)(
160            %Equal.sym(List<&2, U32>, Shape.dimensions(Shape.output_shape(plan)), Shape.dimensions(target), same_dims) :
161              Shape.LT(U32.to_nat(index), Shape.logical_numel(_))
162            inside),
163          {\{Holds: proof => proof} : Holds<{Elementwise.expected(Elementwise.KeepLeft{}, plan, input, Core.Packed{target, zeros}, index) ==
164    Laws.source_value(input, Projection.offset_spec(plan, U32.to_nat(index))) : Maybe<&2, F32>}> -> {Elementwise.expected(Elementwise.KeepLeft{}, plan, input, Core.Packed{target, zeros}, index) ==
165    Laws.source_value(input, Projection.offset_spec(plan, U32.to_nat(index))) : Maybe<&2, F32>}}(
166            mapped_value(plan, index, input, target, zeros, right_dims,
167              ElementwiseProof.indices_success(plan, index,
168                {\{Holds: proof => proof} : Holds<{Nat.is_lt(U32.to_nat(index), Shape.elements(Shape.output_shape(plan))) == True{} : Bool}> ->
169                  {Nat.is_lt(U32.to_nat(index), Shape.elements(Shape.output_shape(plan))) == True{} : Bool}}(
170                  output_inside(plan, target, index, same_dims, inside))))))
171      )}