def finish_mapping

bendtorch-core@0.1.0.0 · src/autograd/gradient_PROOF.bend:505

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/transform_PROOF.bend as TransformProof
import bendtorch-core@0.1.0.0/src/autograd/gradient.bend as Gradient
import bendtorch-core@0.1.0.0/src/autograd/gradient_LAWS.bend as Laws

505def finish_mapping(+total_shape: Shape.Shape, +shape: Shape.Shape, -total: Core.Tensor<total_shape>, -contribution: Core.Dynamic,
506  outcome: Gradient.Outcome,
507  -dimensions: {Shape.dimensions(total_shape) == Shape.dimensions(shape) : List<&2, U32>},
508  -loop: LoopInv(total_shape, total, contribution, Nat.add(Shape.elements(total_shape), 0n), outcome)) ->
509  Holds<Laws.Mapping(Gradient.Held{Core.Packed{total_shape, total}}, contribution, shape, outcome)>:
510  match outcome:
511    case Gradient.Refused{after, returned, error}:
512      Holds{Unit{}}
513    case Gradient.Accumulated{after}:
514      Holds{(
515        Equal.trans(Maybe<&2, List<&2, U32>>, Laws.held_dimensions(after), Some{Shape.dimensions(total_shape)}, Some{Shape.dimensions(shape)},
516          Pair.fst(Dims(total_shape, after), Values(total_shape, total, contribution, Nat.add(Shape.elements(total_shape), 0n), after), loop),
517          %dimensions : {Some{Shape.dimensions(total_shape)} == Some{_} : Maybe<&2, List<&2, U32>>}
518          {==}),
519        index => inside => Equal.trans(Maybe<&2, F32>, Laws.held_observe(after, index),
520          partial(Nat.is_lt(U32.to_nat(index), Nat.add(Shape.elements(total_shape), 0n)),
521            Core.checked_observe(total_shape, total, index), Core.checked_observe_dynamic(contribution, index)),
522          Laws.add_word(Core.checked_observe(total_shape, total, index), Core.checked_observe_dynamic(contribution, index)),
523          Pair.snd(Dims(total_shape, after), Values(total_shape, total, contribution, Nat.add(Shape.elements(total_shape), 0n), after), loop)(index)(
524            %Equal.sym(List<&2, U32>, Shape.dimensions(total_shape), Shape.dimensions(shape), dimensions) :
525              {Nat.is_lt(U32.to_nat(index), Shape.logical_numel(_)) == True{} : Bool}
526            inside),
527          %Equal.sym(Bool, Nat.is_lt(U32.to_nat(index), Nat.add(Shape.elements(total_shape), 0n)), True{},
528            %Equal.sym(Nat, Nat.add(Shape.elements(total_shape), 0n), Shape.elements(total_shape), TransformProof.add_zero_right(Shape.elements(total_shape))) :
529              {Nat.is_lt(U32.to_nat(index), _) == True{} : Bool}
530            %Equal.sym(Nat, Shape.elements(total_shape), Shape.numel(Shape.extents(total_shape)), Shape.elements_exact(total_shape)) :
531              {Nat.is_lt(U32.to_nat(index), _) == True{} : Bool}
532            %Equal.sym(List<&2, U32>, Shape.dimensions(total_shape), Shape.dimensions(shape), dimensions) :
533              {Nat.is_lt(U32.to_nat(index), Shape.logical_numel(_)) == True{} : Bool}
534            inside) :
535            {partial(_, Core.checked_observe(total_shape, total, index), Core.checked_observe_dynamic(contribution, index)) ==
536              Laws.add_word(Core.checked_observe(total_shape, total, index), Core.checked_observe_dynamic(contribution, index)) : Maybe<&2, F32>}
537          {==}))}