def indexed_mapping

bendtorch-core@0.1.0.0 · src/autograd/reduce_PROOF.bend:314

import Base
import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
import bendtorch-core@0.1.0.0/src/storage/main.bend as Storage
import bendtorch-core@0.1.0.0/src/tensor/core.bend as Core
import bendtorch-core@0.1.0.0/src/autograd/reduce.bend as Reduce
import bendtorch-core@0.1.0.0/src/autograd/reduce_LAWS.bend as Laws
import bendtorch-core@0.1.0.0/src/tensor/matmul_PROOF.bend as MatmulProof

314def indexed_mapping(+shape: Shape.Shape, tensor: Core.Tensor<shape>, +target: Shape.Shape, +plan: Shape.BroadcastPlan, output: Core.Tensor<target>,
315  +depth: Nat, +length: Nat, -valid: {Storage.bounds(depth, length) == True{} : Bool},
316  -exact: {length == Shape.logical_numel(Shape.dimensions(target)) : Nat},
317  map: Shape.Index<Shape.extents(Shape.left_shape(plan)), Shape.extents(Shape.right_shape(plan)), Shape.extents(Shape.output_shape(plan))>,
318  -same: {Done{output} == Core.new(target, 0.0) : Result<&2, &1, Core.Error, Core.Tensor<target>>},
319  +left: {Shape.dimensions(Shape.left_shape(plan)) == Shape.dimensions(target) : List<&2, U32>},
320  +right: {Shape.dimensions(Shape.right_shape(plan)) == Shape.dimensions(shape) : List<&2, U32>},
321  +count: {Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan))) == Shape.logical_numel(Shape.dimensions(shape)) : Nat}) ->
322  Holds<Laws.MappingWith(Core.Packed{shape, tensor}, target, Done{plan},
323    Reduce.parallel_indexed(Core.Packed{shape, tensor}, target, plan, output, depth, length, valid, exact, map))>:
324  match map:
325    case Shape.Unindexed{zero}:
326      Holds{({==}, values => listed => index => inside =>
327        %Equal.trans(Nat, 0n, Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan))), Shape.logical_numel(Shape.dimensions(shape)), zero, count) :
328          {Core.checked_observe(target, output, index) == Laws.fiber(plan, values, index, _) : Maybe<&2, F32>}
329        {\{Holds: proof => proof} : Holds<Inv(plan, values, target, 0n, output)> -> Inv(plan, values, target, 0n, output)}(
330          zero_inv(plan, values, target, output, same))(index)(inside))}
331    case Shape.Indexed{rows, same_rows}:
332      MatmulProof.take_lift(shape, tensor,
333        pair => Holds<Laws.MappingWith(Core.Packed{shape, tensor}, target, Done{plan}, Reduce.parallel_rows(plan, target, depth, length, valid, exact, pair, rows))>,
334        kept => copy => taken => rows_mapping(shape, tensor, target, plan, depth, length, valid, exact, kept, copy,
335          taken, rows, same_rows, left, right, count))