def allocated_mapping

bendtorch-core@0.1.0.0 · src/tensor/elementwise_parallel_proof.bend:518

import Base
import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
import bendtorch-core@0.1.0.0/src/shape/nat_proofs.bend as Nats
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/tensor/elementwise.bend as Elementwise
import bendtorch-core@0.1.0.0/src/tensor/elementwise_totality.bend as Totality

518def allocated_mapping(+operation: Elementwise.Operation, +plan: Shape.BroadcastPlan, left: Core.Dynamic, right: Core.Dynamic,
519  result: Result<&2, &1, Core.Error, Core.Tensor<Shape.output_shape(plan)>>,
520  +left_fits: {Totality.dimensions(left) == Shape.dimensions(Shape.left_shape(plan)) : List<&2, U32>},
521  +right_fits: {Totality.dimensions(right) == Shape.dimensions(Shape.right_shape(plan)) : List<&2, U32>}) ->
522  Elementwise.MappingModel<operation, plan, left, right, Elementwise.parallel_allocated(operation, plan, left, right, result)>:
523  match result:
524    case Fail{error}:
525      Elementwise.MappingModel{Unit{}}
526    case Done{Core.TTensor{storage, exact}}:
527      match storage:
528        case Storage.Empty{}:
529          Elementwise.MappingModel{index => bound => Empty.absurd(
530            {Core.observe_dynamic(Core.Packed{Shape.output_shape(plan), Core.TTensor{Storage.Empty{}, exact}}, index) == Elementwise.expected(operation, plan, left, right, index) : Maybe<&2, F32>},
531            Arith.lt_zero(U32.to_nat(index),
532              %Equal.sym(Nat, 0n, Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan))), exact) : {Nat.is_lt(U32.to_nat(index), _) == True{} : Bool}
533              Equal.sym(Bool, True{}, Nat.is_lt(U32.to_nat(index), Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan)))), Nats.lt_true(U32.to_nat(index), Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan))), bound))))}
534        case Storage.Owned{+depth, +length, array, balanced, valid}:
535          output_mapping(operation, plan, depth, length, valid, exact, left, right, Elementwise.take(left), Elementwise.take(right),
536            taken(left), taken(right), left_fits, right_fits)