def decided_lift

bendtorch-core@0.1.0.0 · src/tensor/reshape_contract.bend:42

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/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/lifecycle.bend as Lifecycle
import bendtorch-core@0.1.0.0/src/tensor/transforms.bend as Transforms
import bendtorch-core@0.1.0.0/src/tensor/allocation_PROOF.bend as AllocationProof

42def decided_lift(+input: Shape.Shape, +output: Shape.Shape, -storage: Storage.Storage,
43  -valid: {Lifecycle.logical_length(storage) == Shape.logical_numel(Shape.dimensions(input)) : Nat},
44  decision: Shape.BoolCase<Nat.is_eq(Shape.elements(input), Shape.elements(output))>,
45  counted: {Shape.elements(input) == Shape.elements(output) : Nat},
46  -P: Result<&2, &1, Transforms.Error, Core.Dynamic> -> Type,
47  h: @-fits: {Lifecycle.logical_length(storage) == Shape.logical_numel(Shape.dimensions(output)) : Nat} -> P(Done{Core.Packed{output, Core.TTensor{storage, fits}}})) ->
48  P(Transforms.reshape_decided(input, output, Core.TTensor{storage, valid}, decision)):
49  match decision:
50    case Shape.BFalse{no}:
51      Empty.absurd(P(Transforms.reshape_decided(input, output, Core.TTensor{storage, valid}, Shape.BFalse{no})),
52        Nats.contradiction(Nat.is_eq(Shape.elements(input), Shape.elements(output)),
53          AllocationProof.natural_equal_true(Shape.elements(input), Shape.elements(output), counted), no))
54    case Shape.BTrue{yes}:
55      h(Equal.trans(Nat, Lifecycle.logical_length(storage), Shape.logical_numel(Shape.dimensions(input)), Shape.logical_numel(Shape.dimensions(output)), valid, Transforms.elements_equal(input, output, yes)))