def parallel_mapping

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

import Base
import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
import bendtorch-core@0.1.0.0/src/storage/arith.bend as Arith
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

337def parallel_mapping(+shape: Shape.Shape, tensor: Core.Tensor<shape>, +target: Shape.Shape, +plan: Shape.BroadcastPlan,
338  result: Result<&2, &1, Core.Error, Core.Tensor<target>>,
339  -same: {result == Core.new(target, 0.0) : Result<&2, &1, Core.Error, Core.Tensor<target>>},
340  +left: {Shape.dimensions(Shape.left_shape(plan)) == Shape.dimensions(target) : List<&2, U32>},
341  +right: {Shape.dimensions(Shape.right_shape(plan)) == Shape.dimensions(shape) : List<&2, U32>},
342  +count: {Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan))) == Shape.logical_numel(Shape.dimensions(shape)) : Nat}) ->
343  Holds<Laws.MappingWith(Core.Packed{shape, tensor}, target, Done{plan}, Reduce.parallel_allocated(Core.Packed{shape, tensor}, target, plan, result))>:
344  match result:
345    case Fail{error}:
346      Holds{Unit{}}
347    case Done{Core.TTensor{storage, exact}}:
348      match storage:
349        case Storage.Empty{}:
350          Holds{({==}, values => listed => index => inside =>
351            Empty.absurd({Core.checked_observe(target, Core.TTensor{Storage.Empty{}, exact}, index) ==
352              Laws.fiber(plan, values, index, Shape.logical_numel(Shape.dimensions(shape))) : Maybe<&2, F32>},
353              Arith.lt_zero(U32.to_nat(index),
354                %Equal.sym(Nat, 0n, Shape.logical_numel(Shape.dimensions(target)), exact) : {Nat.is_lt(U32.to_nat(index), _) == True{} : Bool}
355                inside)))}
356        case Storage.Owned{+depth, +length, array, balanced, valid}:
357          indexed_mapping(shape, tensor, target, plan, Core.TTensor{Storage.Owned{depth, length, array, balanced, valid}, exact}, depth, length,
358            valid, exact, Shape.index_map(plan), same, left, right, count)