def sum_dims_decided

bendtorch-core@0.1.0.0 · src/autograd/tensor_backward_PROOF.bend:490

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/storage/arith.bend as Arith
import bendtorch-core@0.1.0.0/src/autograd/tensor_backward.bend as Backward
import bendtorch-core@0.1.0.0/src/autograd/tensor_backward_LAWS.bend as Laws
import bendtorch-core@0.1.0.0/src/tensor/contracts/replayed.bend as Replayed
import bendtorch-core@0.1.0.0/src/tensor/contracts/broadcast.bend as PointwiseBroadcast
import bendtorch-core@0.1.0.0/src/tensor/transforms.bend as Transforms
import bendtorch-core@0.1.0.0/src/tensor/lifecycle.bend as Lifecycle
import bendtorch-core@0.1.0.0/src/storage/LAWS.bend as StorageLaws

490def sum_dims_decided(+a: Nat, +shape: Shape.Shape, +input: Shape.Shape, +kept: Shape.Shape, +plan: Shape.BroadcastPlan,
491  planned: {Shape.broadcast_checked(kept, input) == Done{plan} : Result<&2, &2, Shape.Error, Shape.BroadcastPlan>},
492  same: {Shape.dimensions(Shape.output_shape(plan)) == Shape.dimensions(input) : List<&2, U32>},
493  counted: {Nat.is_eq(Shape.elements(shape), Shape.elements(kept)) == True{} : Bool},
494  saved: Backward.Saved, +fill: F32, +ops: List<&2, StorageLaws.Op>,
495  -fits: {Lifecycle.logical_length(Replayed.storage(kept, fill, ops)) == Shape.logical_numel(Shape.dimensions(shape)) : Nat},
496  decision: Shape.BoolCase<Nat.is_eq(Shape.elements(shape), Shape.elements(kept))>) ->
497  Laws.SoleContract<a, input, index => Replayed.value(kept, fill, ops, U32.from_nat(PointwiseBroadcast.left_map(plan, index))),
498    Backward.unsqueezed(a, input, kept, shape, Core.TTensor{Replayed.storage(kept, fill, ops), fits}, decision)>:
499  match decision:
500    case Shape.BFalse{equal}:
501      Empty.absurd(Laws.SoleContract<a, input, index => Replayed.value(kept, fill, ops, U32.from_nat(PointwiseBroadcast.left_map(plan, index))),
502        Backward.unsqueezed(a, input, kept, shape, Core.TTensor{Replayed.storage(kept, fill, ops), fits}, Shape.BFalse{equal})>,
503        Arith.clash(Equal.trans(Bool, False{}, Nat.is_eq(Shape.elements(shape), Shape.elements(kept)), True{}, equal, counted)))
504    case Shape.BTrue{equal}:
505      restored_contract(a, input, index => Replayed.value(kept, fill, ops, U32.from_nat(PointwiseBroadcast.left_map(plan, index))), shape,
506        Backward.vjp_sum(a, input, Transforms.reshape_equal(shape, kept, Core.TTensor{Replayed.storage(kept, fill, ops), fits},
507          Transforms.elements_equal(shape, kept, equal))),
508        Laws.tensor_sum_vjp_pointwise(a, kept, input, plan, planned, same, saved, fill, ops,
509          Equal.trans(Nat, Lifecycle.logical_length(Replayed.storage(kept, fill, ops)), Shape.logical_numel(Shape.dimensions(shape)),
510            Shape.logical_numel(Shape.dimensions(kept)), fits, Transforms.elements_equal(shape, kept, equal))))