src/tensor/permute_values.bend

bendtorch-core@0.1.0.0 · 576 lines · 39 definitions

1import Base
2import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
3import bendtorch-core@0.1.0.0/src/storage/main.bend as Storage
4import bendtorch-core@0.1.0.0/src/storage/LAWS.bend as StorageLaws
5import bendtorch-core@0.1.0.0/src/storage/arith.bend as Arith
6import bendtorch-core@0.1.0.0/src/tensor/core.bend as Core
7import bendtorch-core@0.1.0.0/src/tensor/lifecycle.bend as Lifecycle
8import bendtorch-core@0.1.0.0/src/tensor/layout.bend as Layout
9import bendtorch-core@0.1.0.0/src/tensor/axis_spec.bend as Axis
10import bendtorch-core@0.1.0.0/src/tensor/transforms.bend as Transforms
11import bendtorch-core@0.1.0.0/src/tensor/transpose_spec.bend as Transpose
12import bendtorch-core@0.1.0.0/src/tensor/transpose_PROOF.bend as TransposeProof
13import bendtorch-core@0.1.0.0/src/tensor/permutation_mapping.bend as Mapping
14import bendtorch-core@0.1.0.0/src/tensor/word_conversion.bend as Conversion
15import bendtorch-core@0.1.0.0/src/tensor/from_values_contract.bend as FromValues
16import bendtorch-core@0.1.0.0/src/tensor/contracts/replayed.bend as Replayed
17import bendtorch-core@0.1.0.0/src/tensor/contracts/pointwise_proof.bend as PointwiseProof
18import bendtorch-core@0.1.0.0/src/tensor/index_map.bend as IndexMap
19import bendtorch-core@0.1.0.0/src/tensor/observation.bend as Values
20import bendtorch-core@0.1.0.0/src/tensor/transform_observation.bend as Observation
21import bendtorch-core@0.1.0.0/src/tensor/permute_values_records.bend as Records
22import bendtorch-core@0.1.0.0/src/tensor/permute_build_proof.bend as PermuteProof
23import bendtorch-core@0.1.0.0/src/tensor/contracts/build_facts.bend as Facts
24import bendtorch-core@0.1.0.0/src/tensor/contracts/words_proof.bend as Words
25import bendtorch-core@0.1.0.0/src/tensor/contracts/LAWS.bend as Contracts
26import bendtorch-core@0.1.0.0/src/tensor/contracts/binary_storage_proof.bend as Owned
27import bendtorch-core@0.1.0.0/src/tensor/read_contract.bend as Reads
28import bendtorch-core@0.1.0.0/src/tensor/permute_build.bend as Build
29import bendtorch-core@0.1.0.0/src/storage/model.bend as Model
30import bendtorch-core@0.1.0.0/src/storage/array.bend as ArrayCore
31import bendtorch-core@0.1.0.0/src/storage/api.bend as Api
32import bendtorch-core@0.1.0.0/src/storage/refinement.bend as Refinement
33import bendtorch-core@0.1.0.0/src/storage/certificates.bend as Certificates
34
35# The permute copy as an index-map instance: output index i reads the input
36# at `Values.permute_source(input, axes, i)`, the row-major index of the
37# input coordinate that `axes` sends to output coordinate i.
38
39type Request is Data:
40  Request{input: Shape.Shape, axes: List<&2, U32>}
41
42def source_of(request: Request, index: Nat) -> Nat:
43  Request{+input, +axes} = request
44  Values.permute_source_index(input, axes, index)
45
46# The parallel copy hands the source back as the First half of its clone.
47def kept_source(+input: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>) -> Storage.Storage:
48  Replayed.storage(input, fill, Replayed.snoc(ops, StorageLaws.First{}))
49
50def written(+input: Shape.Shape, +axes: List<&2, U32>, +plan: Layout.Plan<input, axes>, +fill: F32, +ops: List<&2, StorageLaws.Op>, +count: Nat) -> Storage.Storage:
51  IndexMap.target(~Request, ~source_of, count, Request{input, axes}, input, fill, ops, Layout.output(input, axes, plan))
52
53def Final(+input: Shape.Shape, +axes: List<&2, U32>, +plan: Layout.Plan<input, axes>, +fill: F32, +ops: List<&2, StorageLaws.Op>, +count: Nat,
54  -Q: Core.Tensor<input> & Result<&2, &1, Transforms.Error, Core.Tensor<Layout.output(input, axes, plan)>> -> Type) -> Type:
55  @-source_fits: {Lifecycle.logical_length(kept_source(input, fill, ops)) == Shape.logical_numel(Shape.dimensions(input)) : Nat} ->
56  @-target_fits: {Lifecycle.logical_length(written(input, axes, plan, fill, ops, count)) ==
57    Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan))) : Nat} ->
58    Q((Core.TTensor{kept_source(input, fill, ops), source_fits}, Done{Core.TTensor{written(input, axes, plan, fill, ops, count), target_fits}}))
59
60def half_keeps(+shape: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>, +index: U32) ->
61  {Replayed.value(shape, fill, Replayed.snoc(ops, StorageLaws.First{}), index) == Replayed.value(shape, fill, ops, index) : F32}:
62  Replayed.expected_snoc(ops, StorageLaws.First{}, U32.to_nat(Shape.count(shape)), index, fill)
63
64
65def words_snoc(+shape: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>, +previous: Nat, +word: F32, +other: F32,
66  +before: List<&2, F32>, rest: {Contracts.words(shape, fill, ops, previous) == before : List<&2, F32>}, last: {word == other : F32}) ->
67  {List.append(&2, F32, Contracts.words(shape, fill, ops, previous), [word]) == List.append(&2, F32, before, [other]) : List<&2, F32>}:
68  Equal.trans(List<&2, F32>, List.append(&2, F32, Contracts.words(shape, fill, ops, previous), [word]),
69    List.append(&2, F32, before, [word]), List.append(&2, F32, before, [other]),
70    Equal.cong(List<&2, F32>, List<&2, F32>, values => List.append(&2, F32, values, [word]), Contracts.words(shape, fill, ops, previous), before, rest),
71    Equal.cong(F32, List<&2, F32>, value => List.append(&2, F32, before, [value]), word, other, last))
72
73# The First half of a clone lists the same words.
74
75# The First half of a clone lists the same words.
76def words_first(+shape: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>, +count: Nat) ->
77  {Contracts.words(shape, fill, Replayed.snoc(ops, StorageLaws.First{}), count) == Contracts.words(shape, fill, ops, count) : List<&2, F32>}:
78  match count:
79    case 0n:
80      {==}
81    case 1n+previous:
82      +position = U32.from_nat(previous)
83      words_snoc(shape, fill, Replayed.snoc(ops, StorageLaws.First{}), previous,
84        Replayed.value(shape, fill, Replayed.snoc(ops, StorageLaws.First{}), position), Replayed.value(shape, fill, ops, position),
85        Contracts.words(shape, fill, ops, previous), words_first(shape, fill, ops, previous), half_keeps(shape, fill, ops, position))
86
87# `take` of an owned tensor hands back the First half of its clone and the
88# Second half's array.
89
90# The copy loop's write history, replayed on the owned storage that
91# `created` makes, is the owned storage that holds the output spec.
92def chain(+input: Shape.Shape, +axes: List<&2, U32>, +fill: F32, +ops: List<&2, StorageLaws.Op>, +d: Nat, +n: Nat, -f: {Storage.bounds(d, n) == True{} : Bool}, +fits: {Storage.bounds(d, n) == True{} : Bool},
93  +limit: U32, +n_limit: {Nat.is_le(n, U32.to_nat(limit)) == True{} : Bool},
94  +k: Nat, +room: {Nat.is_le(k, n) == True{} : Bool}) ->
95  {StorageLaws.replay(IndexMap.writes(~Request, ~source_of, k, Request{input, axes}, input, fill, ops), Storage.Owned{d, n, Array.new(F32, d, 0.0), ArrayCore.new_balanced(F32, d, 0.0), f}) == Storage.Owned{d, n, PermuteProof.spec(input, axes, fill, ops, k, d, 0n), PermuteProof.spec_balanced(input, axes, fill, ops, k, d, 0n), f} : Storage.Storage}:
96  match k:
97    case 0n:
98      Owned.owned_eq(d, n, Array.new(F32, d, 0.0), PermuteProof.spec(input, axes, fill, ops, 0n, d, 0n), PermuteProof.new_spec(input, axes, fill, ops, d, 0n), ArrayCore.new_balanced(F32, d, 0.0), f)(
99        PermuteProof.spec_balanced(input, axes, fill, ops, 0n, d, 0n))
100    case 1n+j:
101      +index = U32.from_nat(j)
102      +value = PermuteProof.word(input, axes, fill, ops, j)
103      +below_n = Conversion.le_successor_lt(j, n, room)
104      +exact = Conversion.natural_word_exact(j, limit, Facts.le_trans(j, n, U32.to_nat(limit), Conversion.lt_le(j, n, below_n), n_limit))
105      +inside = Equal.trans(Bool, Nat.is_lt(U32.to_nat(index), n), Nat.is_lt(j, n), True{},
106        Equal.cong(Nat, Bool, x => Nat.is_lt(x, n), U32.to_nat(index), j, exact), below_n)
107      +offset = Equal.trans(Nat, Nat.add(Model.offset(d, Refinement.route(d, index)), 0n), U32.to_nat(index), j,
108        Facts.output_position(d, n, index, fits, inside), exact)
109      +written = Equal.trans(Array<F32>, Array.set(F32, PermuteProof.spec(input, axes, fill, ops, j, d, 0n), index, value), ArrayCore.write_path(F32, d, PermuteProof.spec(input, axes, fill, ops, j, d, 0n), Refinement.route(d, index), value),
110        PermuteProof.spec(input, axes, fill, ops, 1n+j, d, 0n), Refinement.law_set_refines(F32, d, PermuteProof.spec(input, axes, fill, ops, j, d, 0n), index, value, PermuteProof.spec_balanced(input, axes, fill, ops, j, d, 0n)),
111        PermuteProof.write_spec(input, axes, fill, ops, j, d, 0n, Refinement.route(d, index), offset))
112      Equal.trans(Storage.Storage, StorageLaws.replay(IndexMap.writes(~Request, ~source_of, 1n+j, Request{input, axes}, input, fill, ops), Storage.Owned{d, n, Array.new(F32, d, 0.0), ArrayCore.new_balanced(F32, d, 0.0), f}),
113        StorageLaws.step(StorageLaws.replay(IndexMap.writes(~Request, ~source_of, j, Request{input, axes}, input, fill, ops), Storage.Owned{d, n, Array.new(F32, d, 0.0), ArrayCore.new_balanced(F32, d, 0.0), f}), StorageLaws.Write{index, value}),
114        Storage.Owned{d, n, PermuteProof.spec(input, axes, fill, ops, 1n+j, d, 0n), PermuteProof.spec_balanced(input, axes, fill, ops, 1n+j, d, 0n), f},
115        FromValues.replay_snoc(IndexMap.writes(~Request, ~source_of, j, Request{input, axes}, input, fill, ops), StorageLaws.Write{index, value}, Storage.Owned{d, n, Array.new(F32, d, 0.0), ArrayCore.new_balanced(F32, d, 0.0), f}),
116        %Equal.sym(Storage.Storage, StorageLaws.replay(IndexMap.writes(~Request, ~source_of, j, Request{input, axes}, input, fill, ops), Storage.Owned{d, n, Array.new(F32, d, 0.0), ArrayCore.new_balanced(F32, d, 0.0), f}), Storage.Owned{d, n, PermuteProof.spec(input, axes, fill, ops, j, d, 0n), PermuteProof.spec_balanced(input, axes, fill, ops, j, d, 0n), f},
117          chain(input, axes, fill, ops, d, n, f, fits, limit, n_limit, j, Conversion.le_predecessor(j, n, room))) :
118          {StorageLaws.step(_, StorageLaws.Write{index, value}) == Storage.Owned{d, n, PermuteProof.spec(input, axes, fill, ops, 1n+j, d, 0n), PermuteProof.spec_balanced(input, axes, fill, ops, 1n+j, d, 0n), f} : Storage.Storage}
119        %Equal.sym(Storage.Storage & Result<&2, &2, Storage.Error, Unit>,
120          Storage.set_checked(d, n, PermuteProof.spec(input, axes, fill, ops, j, d, 0n), index, value, Nat.is_lt(U32.to_nat(index), n), {==}, PermuteProof.spec_balanced(input, axes, fill, ops, j, d, 0n), f),
121          (Storage.Owned{d, n, Array.set(F32, PermuteProof.spec(input, axes, fill, ops, j, d, 0n), index, value),
122            Certificates.set_result(F32, d, PermuteProof.spec(input, axes, fill, ops, j, d, 0n), index, value, Array.set(F32, PermuteProof.spec(input, axes, fill, ops, j, d, 0n), index, value), {==}, PermuteProof.spec_balanced(input, axes, fill, ops, j, d, 0n)), f}, Done{Unit{}}),
123          Api.set_checked_hit(d, n, PermuteProof.spec(input, axes, fill, ops, j, d, 0n), index, value, Nat.is_lt(U32.to_nat(index), n), {==}, PermuteProof.spec_balanced(input, axes, fill, ops, j, d, 0n), f,
124            Equal.sym(Bool, Nat.is_lt(U32.to_nat(index), n), True{}, inside), fits, inside)) :
125          {StorageLaws.owner(Result<&2, &2, Storage.Error, Unit>, _) == Storage.Owned{d, n, PermuteProof.spec(input, axes, fill, ops, 1n+j, d, 0n), PermuteProof.spec_balanced(input, axes, fill, ops, 1n+j, d, 0n), f} : Storage.Storage}
126        Owned.owned_eq(d, n, Array.set(F32, PermuteProof.spec(input, axes, fill, ops, j, d, 0n), index, value), PermuteProof.spec(input, axes, fill, ops, 1n+j, d, 0n), written,
127          Certificates.set_result(F32, d, PermuteProof.spec(input, axes, fill, ops, j, d, 0n), index, value, Array.set(F32, PermuteProof.spec(input, axes, fill, ops, j, d, 0n), index, value), {==}, PermuteProof.spec_balanced(input, axes, fill, ops, j, d, 0n)), f)(
128          PermuteProof.spec_balanced(input, axes, fill, ops, 1n+j, d, 0n)))
129
130
131# The output count of a permute is its input count.
132def count_same(+input: Shape.Shape, +axes: List<&2, U32>, +plan: Layout.Plan<input, axes>) -> {U32.to_nat(Shape.count(Layout.output(input, axes, plan))) == U32.to_nat(Shape.count(input)) : Nat}:
133  Equal.trans(Nat, U32.to_nat(Shape.count(Layout.output(input, axes, plan))), Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan))), U32.to_nat(Shape.count(input)), Reads.count_exact(Layout.output(input, axes, plan)),
134    Equal.trans(Nat, Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan))), Shape.logical_numel(Shape.dimensions(input)), U32.to_nat(Shape.count(input)), PermuteProof.plan_numel(input, axes, plan), Equal.sym(Nat, U32.to_nat(Shape.count(input)), Shape.logical_numel(Shape.dimensions(input)), Reads.count_exact(input))))
135
136def Packed(+input: Shape.Shape, +axes: List<&2, U32>, +plan: Layout.Plan<input, axes>, Q: Core.Tensor<input> & Result<&2, &1, Transforms.Error, Core.Dynamic> -> Type,
137  kept: Storage.Storage, -kept_fits: {Lifecycle.logical_length(kept) == Shape.logical_numel(Shape.dimensions(input)) : Nat},
138  made: Storage.Storage, -made_fits: {Lifecycle.logical_length(made) == Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan))) : Nat}) -> Type:
139  Q(Transforms.copy_pack(input, axes, plan, (Core.TTensor{kept, kept_fits}, Done{Core.TTensor{made, made_fits}})))
140
141def Kept(+input: Shape.Shape, +axes: List<&2, U32>, +plan: Layout.Plan<input, axes>, Q: Core.Tensor<input> & Result<&2, &1, Transforms.Error, Core.Dynamic> -> Type,
142  source: Storage.Storage, made: Storage.Storage) -> Type:
143  @-kf: {Lifecycle.logical_length(StorageLaws.step(source, StorageLaws.First{})) == Shape.logical_numel(Shape.dimensions(input)) : Nat} ->
144  @-ko: {Lifecycle.logical_length(made) == Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan))) : Nat} ->
145    Packed(input, axes, plan, Q, StorageLaws.step(source, StorageLaws.First{}), kf, made, ko)
146
147def cloned_lift(+input: Shape.Shape, +axes: List<&2, U32>, +plan: Layout.Plan<input, axes>, +fill: F32, +ops: List<&2, StorageLaws.Op>,
148  +d: Nat, -f: {Storage.bounds(d, U32.to_nat(Shape.count(Layout.output(input, axes, plan)))) == True{} : Bool}, -fo: {U32.to_nat(Shape.count(Layout.output(input, axes, plan))) == Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan))) : Nat},
149  +ds: Nat, -b: Model.Balanced(F32, ds, Api.arr(Owned.norm(U32.to_nat(Shape.count(input)), ops, Nil{}), ds, fill)), -f2: {Storage.bounds(ds, U32.to_nat(Shape.count(input))) == True{} : Bool},
150  pair: Array<F32> & Array<F32>, -same: {pair == Array.clone(F32, Api.arr(Owned.norm(U32.to_nat(Shape.count(input)), ops, Nil{}), ds, fill)) : Array<F32> & Array<F32>},
151  -sf: {Lifecycle.logical_length(Storage.Owned{ds, U32.to_nat(Shape.count(input)), Api.arr(Owned.norm(U32.to_nat(Shape.count(input)), ops, Nil{}), ds, fill), b, f2}) == Shape.logical_numel(Shape.dimensions(input)) : Nat},
152  -Q: Core.Tensor<input> & Result<&2, &1, Transforms.Error, Core.Dynamic> -> Type,
153  hyp: @-kf: {Lifecycle.logical_length(StorageLaws.owner(Storage.Storage, Storage.finish_clone(ds, U32.to_nat(Shape.count(input)), Api.arr(Owned.norm(U32.to_nat(Shape.count(input)), ops, Nil{}), ds, fill), pair, same, b, f2))) == Shape.logical_numel(Shape.dimensions(input)) : Nat} ->
154    Packed(input, axes, plan, Q, StorageLaws.owner(Storage.Storage, Storage.finish_clone(ds, U32.to_nat(Shape.count(input)), Api.arr(Owned.norm(U32.to_nat(Shape.count(input)), ops, Nil{}), ds, fill), pair, same, b, f2)), kf,
155      Storage.Owned{d, U32.to_nat(Shape.count(Layout.output(input, axes, plan))), Build.build(d, Build.Request{input, axes}, U32.to_nat(Shape.count(Layout.output(input, axes, plan))), Pair.snd(Array<F32>, Array<F32>, pair)),
156        Build.balanced(d, Build.Request{input, axes}, U32.to_nat(Shape.count(Layout.output(input, axes, plan))), Pair.snd(Array<F32>, Array<F32>, pair)), f}, fo)) ->
157  Q(Transforms.copy_pack(input, axes, plan, Transforms.built_cloned(input, axes, plan, d, U32.to_nat(Shape.count(Layout.output(input, axes, plan))), f, fo, ds, U32.to_nat(Shape.count(input)), Api.arr(Owned.norm(U32.to_nat(Shape.count(input)), ops, Nil{}), ds, fill), pair, same, b, f2, sf))):
158  (kept, copy) = pair
159  hyp(sf)
160
161# The copy loop's write history, replayed on the storage `Core.new`
162# allocates, is the build's output on the Second half of the source's clone.
163def made(+input: Shape.Shape, +axes: List<&2, U32>, +plan: Layout.Plan<input, axes>, +fill: F32, +ops: List<&2, StorageLaws.Op>,
164  +d: Nat, +fits: {Storage.bounds(d, U32.to_nat(Shape.count(Layout.output(input, axes, plan)))) == True{} : Bool}, -f: {Storage.bounds(d, U32.to_nat(Shape.count(Layout.output(input, axes, plan)))) == True{} : Bool},
165  +ds: Nat, +fits_s: {Storage.bounds(ds, U32.to_nat(Shape.count(input))) == True{} : Bool}) ->
166  {StorageLaws.replay(IndexMap.writes(~Request, ~source_of, Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan))), Request{input, axes}, input, fill, ops), Storage.Owned{d, U32.to_nat(Shape.count(Layout.output(input, axes, plan))), Array.new(F32, d, 0.0), ArrayCore.new_balanced(F32, d, 0.0), f}) == Storage.Owned{d, U32.to_nat(Shape.count(Layout.output(input, axes, plan))), Build.build(d, Build.Request{input, axes}, U32.to_nat(Shape.count(Layout.output(input, axes, plan))), Api.arr(Api.Right{} <> Owned.norm(U32.to_nat(Shape.count(input)), ops, Nil{}), ds, fill)),
167    Build.balanced(d, Build.Request{input, axes}, U32.to_nat(Shape.count(Layout.output(input, axes, plan))), Api.arr(Api.Right{} <> Owned.norm(U32.to_nat(Shape.count(input)), ops, Nil{}), ds, fill)), f} : Storage.Storage}:
168  +count_exact = Reads.count_exact(Layout.output(input, axes, plan))
169  Equal.trans(Storage.Storage, StorageLaws.replay(IndexMap.writes(~Request, ~source_of, Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan))), Request{input, axes}, input, fill, ops), Storage.Owned{d, U32.to_nat(Shape.count(Layout.output(input, axes, plan))), Array.new(F32, d, 0.0), ArrayCore.new_balanced(F32, d, 0.0), f}),
170    StorageLaws.replay(IndexMap.writes(~Request, ~source_of, U32.to_nat(Shape.count(Layout.output(input, axes, plan))), Request{input, axes}, input, fill, ops), Storage.Owned{d, U32.to_nat(Shape.count(Layout.output(input, axes, plan))), Array.new(F32, d, 0.0), ArrayCore.new_balanced(F32, d, 0.0), f}),
171    Storage.Owned{d, U32.to_nat(Shape.count(Layout.output(input, axes, plan))), Build.build(d, Build.Request{input, axes}, U32.to_nat(Shape.count(Layout.output(input, axes, plan))), Api.arr(Api.Right{} <> Owned.norm(U32.to_nat(Shape.count(input)), ops, Nil{}), ds, fill)), Build.balanced(d, Build.Request{input, axes}, U32.to_nat(Shape.count(Layout.output(input, axes, plan))), Api.arr(Api.Right{} <> Owned.norm(U32.to_nat(Shape.count(input)), ops, Nil{}), ds, fill)), f},
172    Equal.cong(Nat, Storage.Storage, count => StorageLaws.replay(IndexMap.writes(~Request, ~source_of, count, Request{input, axes}, input, fill, ops), Storage.Owned{d, U32.to_nat(Shape.count(Layout.output(input, axes, plan))), Array.new(F32, d, 0.0), ArrayCore.new_balanced(F32, d, 0.0), f}),
173      Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan))), U32.to_nat(Shape.count(Layout.output(input, axes, plan))), Equal.sym(Nat, U32.to_nat(Shape.count(Layout.output(input, axes, plan))), Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan))), count_exact)),
174    Equal.trans(Storage.Storage, StorageLaws.replay(IndexMap.writes(~Request, ~source_of, U32.to_nat(Shape.count(Layout.output(input, axes, plan))), Request{input, axes}, input, fill, ops), Storage.Owned{d, U32.to_nat(Shape.count(Layout.output(input, axes, plan))), Array.new(F32, d, 0.0), ArrayCore.new_balanced(F32, d, 0.0), f}),
175      Storage.Owned{d, U32.to_nat(Shape.count(Layout.output(input, axes, plan))), PermuteProof.spec(input, axes, fill, ops, U32.to_nat(Shape.count(Layout.output(input, axes, plan))), d, 0n), PermuteProof.spec_balanced(input, axes, fill, ops, U32.to_nat(Shape.count(Layout.output(input, axes, plan))), d, 0n), f},
176      Storage.Owned{d, U32.to_nat(Shape.count(Layout.output(input, axes, plan))), Build.build(d, Build.Request{input, axes}, U32.to_nat(Shape.count(Layout.output(input, axes, plan))), Api.arr(Api.Right{} <> Owned.norm(U32.to_nat(Shape.count(input)), ops, Nil{}), ds, fill)), Build.balanced(d, Build.Request{input, axes}, U32.to_nat(Shape.count(Layout.output(input, axes, plan))), Api.arr(Api.Right{} <> Owned.norm(U32.to_nat(Shape.count(input)), ops, Nil{}), ds, fill)), f},
177      chain(input, axes, fill, ops, d, U32.to_nat(Shape.count(Layout.output(input, axes, plan))), f, fits, Shape.count(Layout.output(input, axes, plan)), Conversion.le_reflexive(U32.to_nat(Shape.count(Layout.output(input, axes, plan)))), U32.to_nat(Shape.count(Layout.output(input, axes, plan))), Conversion.le_reflexive(U32.to_nat(Shape.count(Layout.output(input, axes, plan))))),
178      Owned.owned_eq(d, U32.to_nat(Shape.count(Layout.output(input, axes, plan))), PermuteProof.spec(input, axes, fill, ops, U32.to_nat(Shape.count(Layout.output(input, axes, plan))), d, 0n), Build.build(d, Build.Request{input, axes}, U32.to_nat(Shape.count(Layout.output(input, axes, plan))), Api.arr(Api.Right{} <> Owned.norm(U32.to_nat(Shape.count(input)), ops, Nil{}), ds, fill)),
179        Equal.sym(Array<F32>, Build.build(d, Build.Request{input, axes}, U32.to_nat(Shape.count(Layout.output(input, axes, plan))), Api.arr(Api.Right{} <> Owned.norm(U32.to_nat(Shape.count(input)), ops, Nil{}), ds, fill)), PermuteProof.spec(input, axes, fill, ops, U32.to_nat(Shape.count(Layout.output(input, axes, plan))), d, 0n),
180          PermuteProof.build_eq(input, axes, fill, ops, plan, ds, fits_s, U32.to_nat(Shape.count(Layout.output(input, axes, plan))), count_exact, d, Build.splits(d), 0n)),
181        PermuteProof.spec_balanced(input, axes, fill, ops, U32.to_nat(Shape.count(Layout.output(input, axes, plan))), d, 0n), f)(Build.balanced(d, Build.Request{input, axes}, U32.to_nat(Shape.count(Layout.output(input, axes, plan))), Api.arr(Api.Right{} <> Owned.norm(U32.to_nat(Shape.count(input)), ops, Nil{}), ds, fill)))))
182
183# With no output cell the input has none either: both are empty, and the
184# copy hands the source back as it is.
185def empty_lift(+input: Shape.Shape, +axes: List<&2, U32>, +plan: Layout.Plan<input, axes>, +fill: F32, +ops: List<&2, StorageLaws.Op>,
186  +zero: {U32.to_nat(Shape.count(Layout.output(input, axes, plan))) == 0n : Nat}, -Q: Core.Tensor<input> & Result<&2, &1, Transforms.Error, Core.Dynamic> -> Type,
187  -fo: {Lifecycle.logical_length(Storage.Empty{}) == Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan))) : Nat}) ->
188  @-sf: {Lifecycle.logical_length(Replayed.storage(input, fill, ops)) == Shape.logical_numel(Shape.dimensions(input)) : Nat} -> Kept(input, axes, plan, Q, Replayed.storage(input, fill, ops), StorageLaws.replay(IndexMap.writes(~Request, ~source_of, Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan))), Request{input, axes}, input, fill, ops), Storage.Empty{})) ->
189    Packed(input, axes, plan, Q, Replayed.storage(input, fill, ops), sf, Storage.Empty{}, fo):
190  Owned.created_lift(Shape.count(input), fill, Conversion.shape_count_within(input),
191    s => (@-sf: {Lifecycle.logical_length(StorageLaws.replay(ops, s)) == Shape.logical_numel(Shape.dimensions(input)) : Nat} -> Kept(input, axes, plan, Q, StorageLaws.replay(ops, s), StorageLaws.replay(IndexMap.writes(~Request, ~source_of, Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan))), Request{input, axes}, input, fill, ops), Storage.Empty{})) ->
192      Packed(input, axes, plan, Q, StorageLaws.replay(ops, s), sf, Storage.Empty{}, fo)),
193    nothing => sf => hyp =>
194      Facts.storage_shift(input, StorageLaws.step(StorageLaws.replay(ops, Storage.Empty{}), StorageLaws.First{}), StorageLaws.replay(ops, Storage.Empty{}),
195        Equal.trans(Storage.Storage, StorageLaws.step(StorageLaws.replay(ops, Storage.Empty{}), StorageLaws.First{}), Storage.Empty{}, StorageLaws.replay(ops, Storage.Empty{}),
196          Equal.cong(Storage.Storage, Storage.Storage, st => StorageLaws.step(st, StorageLaws.First{}), StorageLaws.replay(ops, Storage.Empty{}), Storage.Empty{}, Owned.replay_empty(ops)),
197          Equal.sym(Storage.Storage, StorageLaws.replay(ops, Storage.Empty{}), Storage.Empty{}, Owned.replay_empty(ops))),
198        st => k => Packed(input, axes, plan, Q, st, k, Storage.Empty{}, fo),
199        k => Facts.storage_shift(Layout.output(input, axes, plan), StorageLaws.replay(IndexMap.writes(~Request, ~source_of, Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan))), Request{input, axes}, input, fill, ops), Storage.Empty{}), Storage.Empty{}, Owned.replay_empty(IndexMap.writes(~Request, ~source_of, Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan))), Request{input, axes}, input, fill, ops)),
200          st => ko => Packed(input, axes, plan, Q, StorageLaws.step(StorageLaws.replay(ops, Storage.Empty{}), StorageLaws.First{}), k, st, ko),
201          ko => hyp(k, ko))(fo))(sf),
202    d => fits => f => Empty.absurd(@-sf: {Lifecycle.logical_length(StorageLaws.replay(ops, Storage.Owned{d, U32.to_nat(Shape.count(input)), Array.new(F32, d, fill), ArrayCore.new_balanced(F32, d, fill), f})) == Shape.logical_numel(Shape.dimensions(input)) : Nat} ->
203      Kept(input, axes, plan, Q, StorageLaws.replay(ops, Storage.Owned{d, U32.to_nat(Shape.count(input)), Array.new(F32, d, fill), ArrayCore.new_balanced(F32, d, fill), f}), StorageLaws.replay(IndexMap.writes(~Request, ~source_of, Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan))), Request{input, axes}, input, fill, ops), Storage.Empty{})) ->
204        Packed(input, axes, plan, Q, StorageLaws.replay(ops, Storage.Owned{d, U32.to_nat(Shape.count(input)), Array.new(F32, d, fill), ArrayCore.new_balanced(F32, d, fill), f}), sf, Storage.Empty{}, fo),
205      Arith.lt_zero(0n,
206        %Equal.trans(Nat, U32.to_nat(Shape.count(input)), U32.to_nat(Shape.count(Layout.output(input, axes, plan))), 0n, Equal.sym(Nat, U32.to_nat(Shape.count(Layout.output(input, axes, plan))), U32.to_nat(Shape.count(input)), count_same(input, axes, plan)), zero) : {Nat.is_lt(0n, _) == True{} : Bool}
207        Arith.and_left(Nat.is_lt(0n, U32.to_nat(Shape.count(input))), Nat.is_le(U32.to_nat(Shape.count(input)), Model.capacity(d)),
208          Arith.and_right(Nat.is_le(d, 20n), Bool.and(Nat.is_lt(0n, U32.to_nat(Shape.count(input))), Nat.is_le(U32.to_nat(Shape.count(input)), Model.capacity(d))), fits)))))
209
210def run_lift(+input: Shape.Shape, +axes: List<&2, U32>, +plan: Layout.Plan<input, axes>, +fill: F32, +ops: List<&2, StorageLaws.Op>,
211  -Q: Core.Tensor<input> & Result<&2, &1, Transforms.Error, Core.Dynamic> -> Type,
212  h: Final(input, axes, plan, fill, ops, Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan))), pair => Q(Transforms.copy_pack(input, axes, plan, pair))),
213  -source_fits: {Lifecycle.logical_length(Replayed.storage(input, fill, ops)) == Shape.logical_numel(Shape.dimensions(input)) : Nat}) ->
214  Q(Transforms.copy_allocated(input, axes, plan, Core.TTensor{Replayed.storage(input, fill, ops), source_fits}, Core.new(Layout.output(input, axes, plan), 0.0))):
215  Replayed.new_lift(Layout.output(input, axes, plan), 0.0, result => Q(Transforms.copy_allocated(input, axes, plan, Core.TTensor{Replayed.storage(input, fill, ops), source_fits}, result)),
216    target_fits => Owned.created_lift(Shape.count(Layout.output(input, axes, plan)), 0.0, Conversion.shape_count_within(Layout.output(input, axes, plan)),
217      s => (Kept(input, axes, plan, Q, Replayed.storage(input, fill, ops), StorageLaws.replay(IndexMap.writes(~Request, ~source_of, Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan))), Request{input, axes}, input, fill, ops), s)) -> @-fo: {Lifecycle.logical_length(s) == Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan))) : Nat} ->
218        Q(Transforms.copy_pack(input, axes, plan, Transforms.built(input, axes, plan, Core.TTensor{Replayed.storage(input, fill, ops), source_fits}, Core.TTensor{s, fo})))),
219      zero => hyp => fo => empty_lift(input, axes, plan, fill, ops, zero, Q, fo)(source_fits, hyp),
220      d => fits => f => hyp => fo =>
221        +positive = Arith.and_left(Nat.is_lt(0n, U32.to_nat(Shape.count(Layout.output(input, axes, plan)))), Nat.is_le(U32.to_nat(Shape.count(Layout.output(input, axes, plan))), Model.capacity(d)),
222          Arith.and_right(Nat.is_le(d, 20n), Bool.and(Nat.is_lt(0n, U32.to_nat(Shape.count(Layout.output(input, axes, plan)))), Nat.is_le(U32.to_nat(Shape.count(Layout.output(input, axes, plan))), Model.capacity(d))), fits))
223        Facts.input_lift(input, fill, ops,
224          %count_same(input, axes, plan) : {Nat.is_lt(0n, _) == True{} : Bool}
225          positive,
226          st => (@-sf: {Lifecycle.logical_length(st) == Shape.logical_numel(Shape.dimensions(input)) : Nat} -> Kept(input, axes, plan, Q, st, StorageLaws.replay(IndexMap.writes(~Request, ~source_of, Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan))), Request{input, axes}, input, fill, ops), Storage.Owned{d, U32.to_nat(Shape.count(Layout.output(input, axes, plan))), Array.new(F32, d, 0.0), ArrayCore.new_balanced(F32, d, 0.0), f})) ->
227            Q(Transforms.copy_pack(input, axes, plan, Transforms.built_owned(input, axes, plan, d, U32.to_nat(Shape.count(Layout.output(input, axes, plan))), f, fo, Core.TTensor{st, sf})))),
228          ds => fits_s => b => f2 => sf => hyp2 =>
229            cloned_lift(input, axes, plan, fill, ops, d, f, fo, ds, b, f2, Array.clone(F32, Api.arr(Owned.norm(U32.to_nat(Shape.count(input)), ops, Nil{}), ds, fill)), {==}, sf, Q,
230              kf => Facts.storage_shift(Layout.output(input, axes, plan), StorageLaws.replay(IndexMap.writes(~Request, ~source_of, Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan))), Request{input, axes}, input, fill, ops), Storage.Owned{d, U32.to_nat(Shape.count(Layout.output(input, axes, plan))), Array.new(F32, d, 0.0), ArrayCore.new_balanced(F32, d, 0.0), f}),
231                Storage.Owned{d, U32.to_nat(Shape.count(Layout.output(input, axes, plan))), Build.build(d, Build.Request{input, axes}, U32.to_nat(Shape.count(Layout.output(input, axes, plan))), Api.arr(Api.Right{} <> Owned.norm(U32.to_nat(Shape.count(input)), ops, Nil{}), ds, fill)), Build.balanced(d, Build.Request{input, axes}, U32.to_nat(Shape.count(Layout.output(input, axes, plan))), Api.arr(Api.Right{} <> Owned.norm(U32.to_nat(Shape.count(input)), ops, Nil{}), ds, fill)), f},
232                made(input, axes, plan, fill, ops, d, fits, f, ds, fits_s),
233                st => ko => Packed(input, axes, plan, Q, StorageLaws.step(Storage.Owned{ds, U32.to_nat(Shape.count(input)), Api.arr(Owned.norm(U32.to_nat(Shape.count(input)), ops, Nil{}), ds, fill), b, f2}, StorageLaws.First{}), kf, st, ko),
234                ko => hyp2(kf, ko))(fo)))(source_fits, hyp))(
235      kf => ko => Facts.storage_shift(input, kept_source(input, fill, ops), StorageLaws.step(Replayed.storage(input, fill, ops), StorageLaws.First{}),
236        Replayed.replay_snoc(ops, StorageLaws.First{}, input, fill),
237        st => k => Packed(input, axes, plan, Q, st, k, StorageLaws.replay(IndexMap.writes(~Request, ~source_of, Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan))), Request{input, axes}, input, fill, ops), StorageLaws.created(Shape.count(Layout.output(input, axes, plan)), 0.0)), ko),
238        k => h(k, ko))(kf), target_fits))
239
240def planned_lift(+input: Shape.Shape, +axes: List<&2, U32>, +fill: F32, +ops: List<&2, StorageLaws.Op>,
241  result: Result<&2, &2, Layout.Error, Layout.Plan<input, axes>>,
242  done: {Result.is_done(&2, &2, Layout.Error, Layout.Plan<input, axes>, result) == True{} : Bool},
243  -Q: Core.Tensor<input> & Result<&2, &1, Transforms.Error, Core.Dynamic> -> Type,
244  h: @+plan: Layout.Plan<input, axes> -> Final(input, axes, plan, fill, ops, Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan))),
245    pair => Q(Transforms.copy_pack(input, axes, plan, pair))),
246  -source_fits: {Lifecycle.logical_length(Replayed.storage(input, fill, ops)) == Shape.logical_numel(Shape.dimensions(input)) : Nat}) ->
247  Q(Transforms.copy_planned(input, axes, Core.TTensor{Replayed.storage(input, fill, ops), source_fits}, result)):
248  match result:
249    case Fail{error}:
250      Empty.absurd(Q(Transforms.copy_planned(input, axes, Core.TTensor{Replayed.storage(input, fill, ops), source_fits},
251        Fail{error})), Arith.clash(done))
252    case Done{+plan}:
253      run_lift(input, axes, plan, fill, ops, Q, h(plan), source_fits)
254
255def permute_lift(+input: Shape.Shape, +values: List<&2, F32>, +axes: List<&2, U32>, +valid: {True{} == Layout.bijective(input, axes) : Bool},
256  -Q: Core.Tensor<input> & Result<&2, &1, Transforms.Error, Core.Dynamic> -> Type,
257  h: @+plan: Layout.Plan<input, axes> -> Final(input, axes, plan, 0.0, FromValues.writes(values, 0n), Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan))),
258    pair => Q(Transforms.copy_pack(input, axes, plan, pair))),
259  -source_fits: {Lifecycle.logical_length(Replayed.storage(input, 0.0, FromValues.writes(values, 0n))) == Shape.logical_numel(Shape.dimensions(input)) : Nat}) ->
260  Q(Transforms.permute_copy(input, Core.TTensor{Replayed.storage(input, 0.0, FromValues.writes(values, 0n)), source_fits}, axes)):
261  planned_lift(input, axes, 0.0, FromValues.writes(values, 0n), Layout.permutation(input, axes),
262    Mapping.plan_ranked_success(input, axes,
263      Shape.bool_case(Nat.is_eq(Shape.rank_dimensions(axes), Shape.rank_dimensions(Shape.dimensions(input)))), valid),
264    Q, h, source_fits)
265
266def output_inside(+input: Shape.Shape, +axes: List<&2, U32>, +plan: Layout.Plan<input, axes>, +index: U32,
267  inside: {Nat.is_lt(U32.to_nat(index), Shape.logical_numel(Shape.dimensions(input))) == True{} : Bool}) ->
268  {Nat.is_lt(U32.to_nat(index), Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan)))) == True{} : Bool}:
269  %Equal.sym(Nat, Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan))), Shape.logical_numel(Shape.dimensions(input)),
270    PermuteProof.plan_numel(input, axes, plan)) : {Nat.is_lt(U32.to_nat(index), _) == True{} : Bool}
271  inside
272
273def found_value(+input: Shape.Shape, +axes: List<&2, U32>, +plan: Layout.Plan<input, axes>, +values: List<&2, F32>,
274  +counted: {List.length(&2, F32, values) == Shape.logical_numel(Shape.dimensions(input)) : Nat},
275  +index: U32, +mapping: Layout.Mapping<input, axes, plan, index>) ->
276  {Some{Replayed.value(input, 0.0, FromValues.writes(values, 0n), U32.from_nat(source_of(Request{input, axes}, U32.to_nat(index))))} ==
277    Values.listed(values, Values.permute_source(input, axes, index)) : Maybe<&2, F32>}:
278  +found = Layout.source_index(input, axes, plan, index, mapping)
279  +history = FromValues.writes(values, 0n)
280  %Equal.sym(U32, U32.from_nat(U32.to_nat(index)), index, Conversion.word_natural_exact(index)) :
281    {Some{Replayed.value(input, 0.0, history, U32.from_nat(U32.to_nat(Maybe.default(&2, U32, Values.permute_source(input, axes, _), 0))))} ==
282      Values.listed(values, Values.permute_source(input, axes, index)) : Maybe<&2, F32>}
283  %PermuteProof.mapped_index(input, axes, plan, index, mapping) :
284    {Some{Replayed.value(input, 0.0, history, _)} == Values.listed(values, Values.permute_source(input, axes, index)) : Maybe<&2, F32>}
285  %PermuteProof.reference(input, axes, plan, index, mapping) :
286    {Some{Replayed.value(input, 0.0, history, found)} == Values.listed(values, _) : Maybe<&2, F32>}
287  IndexMap.stored_at(input, values, counted, found, PermuteProof.mapped_bound(input, axes, plan, index, mapping))
288
289def found_mapped(+input: Shape.Shape, +axes: List<&2, U32>, +plan: Layout.Plan<input, axes>, +values: List<&2, F32>,
290  +counted: {List.length(&2, F32, values) == Shape.logical_numel(Shape.dimensions(input)) : Nat}, +index: U32,
291  result: Result<&2, &2, Layout.Error, Layout.Mapping<input, axes, plan, index>>,
292  done: {Result.is_done(&2, &2, Layout.Error, Layout.Mapping<input, axes, plan, index>, result) == True{} : Bool}) ->
293  {Some{Replayed.value(input, 0.0, FromValues.writes(values, 0n), U32.from_nat(source_of(Request{input, axes}, U32.to_nat(index))))} ==
294    Values.listed(values, Values.permute_source(input, axes, index)) : Maybe<&2, F32>}:
295  match result:
296    case Fail{error}:
297      Empty.absurd({Some{Replayed.value(input, 0.0, FromValues.writes(values, 0n), U32.from_nat(source_of(Request{input, axes}, U32.to_nat(index))))} ==
298        Values.listed(values, Values.permute_source(input, axes, index)) : Maybe<&2, F32>}, Arith.clash(done))
299    case Done{+mapping}:
300      found_value(input, axes, plan, values, counted, index, mapping)
301
302def Permuted(+input: Shape.Shape, +values: List<&2, F32>, +axes: List<&2, U32>,
303  -fits: {Lifecycle.logical_length(Replayed.storage(input, 0.0, FromValues.writes(values, 0n))) == Shape.logical_numel(Shape.dimensions(input)) : Nat},
304  Q: Core.Tensor<input> & Result<&2, &1, Transforms.Error, Core.Dynamic> -> Type) -> Type:
305  Q(Transforms.permute_copy(input, Core.TTensor{Replayed.storage(input, 0.0, FromValues.writes(values, 0n)), fits}, axes))
306
307def permuted_code(+input: Shape.Shape, +values: List<&2, F32>, +axes: List<&2, U32>, +valid: {True{} == Layout.bijective(input, axes) : Bool},
308  -fits: {Lifecycle.logical_length(Replayed.storage(input, 0.0, FromValues.writes(values, 0n))) == Shape.logical_numel(Shape.dimensions(input)) : Nat}) ->
309  Permuted(input, values, axes, fits, pair => {Some{Observation.failure_code_pair(input, pair)} == Some{0} : Maybe<&2, U32>}):
310  permute_lift(input, values, axes, valid, pair => {Some{Observation.failure_code_pair(input, pair)} == Some{0} : Maybe<&2, U32>},
311    plan => source_fits => target_fits => {==}, fits)
312
313def permuted_dimensions(+input: Shape.Shape, +values: List<&2, F32>, +axes: List<&2, U32>, +valid: {True{} == Layout.bijective(input, axes) : Bool},
314  -fits: {Lifecycle.logical_length(Replayed.storage(input, 0.0, FromValues.writes(values, 0n))) == Shape.logical_numel(Shape.dimensions(input)) : Nat}) ->
315  Permuted(input, values, axes, fits, pair => {Observation.applied_dimensions_pair(input, pair) ==
316    Some{Axis.order(U32, Shape.dimensions(input), Shape.naturals(axes))} : Maybe<&2, List<&2, U32>>}):
317  permute_lift(input, values, axes, valid, pair => {Observation.applied_dimensions_pair(input, pair) ==
318      Some{Axis.order(U32, Shape.dimensions(input), Shape.naturals(axes))} : Maybe<&2, List<&2, U32>>},
319    plan => source_fits => target_fits =>
320      Equal.cong(List<&2, U32>, Maybe<&2, List<&2, U32>>, dimensions => Some{dimensions},
321        Shape.dimensions(Layout.output(input, axes, plan)), Axis.order(U32, Shape.dimensions(input), Shape.naturals(axes)),
322        Mapping.permutation_dimensions_order(input, axes, plan)), fits)
323
324def permuted_length(+input: Shape.Shape, +values: List<&2, F32>, +axes: List<&2, U32>, +valid: {True{} == Layout.bijective(input, axes) : Bool},
325  -fits: {Lifecycle.logical_length(Replayed.storage(input, 0.0, FromValues.writes(values, 0n))) == Shape.logical_numel(Shape.dimensions(input)) : Nat}) ->
326  Permuted(input, values, axes, fits, pair => {Observation.applied_length_pair(input, pair) ==
327    Some{Shape.logical_numel(Shape.dimensions(input))} : Maybe<&2, Nat>}):
328  permute_lift(input, values, axes, valid, pair => {Observation.applied_length_pair(input, pair) ==
329      Some{Shape.logical_numel(Shape.dimensions(input))} : Maybe<&2, Nat>},
330    plan => source_fits => target_fits =>
331      +output = Layout.output(input, axes, plan)
332      +total = Shape.logical_numel(Shape.dimensions(output))
333      -made = written(input, axes, plan, 0.0, FromValues.writes(values, 0n), total)
334      Equal.cong(Nat, Maybe<&2, Nat>, number => Some{number}, Lifecycle.logical_length(made), Shape.logical_numel(Shape.dimensions(input)),
335        Equal.trans(Nat, Lifecycle.logical_length(made), total, Shape.logical_numel(Shape.dimensions(input)),
336          Replayed.fits(output, 0.0, IndexMap.writes(~Request, ~source_of, total, Request{input, axes}, input, 0.0, FromValues.writes(values, 0n))),
337          PermuteProof.plan_numel(input, axes, plan))), fits)
338
339def permuted_count(+input: Shape.Shape, +values: List<&2, F32>, +axes: List<&2, U32>, +valid: {True{} == Layout.bijective(input, axes) : Bool},
340  -fits: {Lifecycle.logical_length(Replayed.storage(input, 0.0, FromValues.writes(values, 0n))) == Shape.logical_numel(Shape.dimensions(input)) : Nat}) ->
341  Permuted(input, values, axes, fits, pair => {Observation.applied_count_pair(input, pair) == Some{Shape.count(input)} : Maybe<&2, U32>}):
342  permute_lift(input, values, axes, valid, pair => {Observation.applied_count_pair(input, pair) == Some{Shape.count(input)} : Maybe<&2, U32>},
343    plan => source_fits => target_fits =>
344      Equal.cong(Nat, Maybe<&2, U32>, number => Some{U32.from_nat(number)}, Shape.elements(Layout.output(input, axes, plan)), Shape.elements(input),
345        PermuteProof.plan_elements(input, axes, plan)), fits)
346
347def permuted_value(+input: Shape.Shape, +values: List<&2, F32>,
348  +counted: {List.length(&2, F32, values) == Shape.logical_numel(Shape.dimensions(input)) : Nat},
349  +axes: List<&2, U32>, +valid: {True{} == Layout.bijective(input, axes) : Bool}, +index: U32,
350  +inside: {Nat.is_lt(U32.to_nat(index), Shape.logical_numel(Shape.dimensions(input))) == True{} : Bool},
351  -fits: {Lifecycle.logical_length(Replayed.storage(input, 0.0, FromValues.writes(values, 0n))) == Shape.logical_numel(Shape.dimensions(input)) : Nat}) ->
352  Permuted(input, values, axes, fits, pair => {Observation.applied_value_pair(input, index, pair) ==
353    Values.listed(values, Values.permute_source(input, axes, index)) : Maybe<&2, F32>}):
354  permute_lift(input, values, axes, valid, pair => {Observation.applied_value_pair(input, index, pair) ==
355      Values.listed(values, Values.permute_source(input, axes, index)) : Maybe<&2, F32>},
356    plan => source_fits => target_fits =>
357      +output = Layout.output(input, axes, plan)
358      +total = Shape.logical_numel(Shape.dimensions(output))
359      +history = FromValues.writes(values, 0n)
360      +below = output_inside(input, axes, plan, index, inside)
361      Equal.trans(Maybe<&2, F32>,
362        Core.dynamic_read_value(Core.read_dynamic(Core.Packed{output, Core.TTensor{written(input, axes, plan, 0.0, FromValues.writes(values, 0n), total), target_fits}}, index)),
363        Some{Replayed.value(input, 0.0, history, U32.from_nat(source_of(Request{input, axes}, U32.to_nat(index))))},
364        Values.listed(values, Values.permute_source(input, axes, index)),
365        IndexMap.copied_point(~Request, ~source_of, Request{input, axes}, input, 0.0, history, output, target_fits, index, below),
366        found_mapped(input, axes, plan, values, counted, index, Layout.map_index(input, axes, plan, index),
367          PermuteProof.mapping_done(input, axes, plan, index, below))), fits)
368
369def permuted_kept(+input: Shape.Shape, +values: List<&2, F32>,
370  +counted: {List.length(&2, F32, values) == Shape.logical_numel(Shape.dimensions(input)) : Nat},
371  +axes: List<&2, U32>, +valid: {True{} == Layout.bijective(input, axes) : Bool}, +other: U32,
372  +inside: {Nat.is_lt(U32.to_nat(other), Shape.logical_numel(Shape.dimensions(input))) == True{} : Bool},
373  -fits: {Lifecycle.logical_length(Replayed.storage(input, 0.0, FromValues.writes(values, 0n))) == Shape.logical_numel(Shape.dimensions(input)) : Nat}) ->
374  Permuted(input, values, axes, fits, pair => {Observation.kept_value_pair(input, other, pair) ==
375    List.get(&2, F32, values, U32.to_nat(other)) : Maybe<&2, F32>}):
376  permute_lift(input, values, axes, valid, pair => {Observation.kept_value_pair(input, other, pair) ==
377      List.get(&2, F32, values, U32.to_nat(other)) : Maybe<&2, F32>},
378    plan => source_fits => target_fits =>
379      +total = Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan)))
380      +history = FromValues.writes(values, 0n)
381      Equal.trans(Maybe<&2, F32>,
382        Core.read_value(input, Core.read_flat(input, Core.TTensor{kept_source(input, 0.0, FromValues.writes(values, 0n)), source_fits}, other)),
383        Some{Replayed.value(input, 0.0, history, other)}, List.get(&2, F32, values, U32.to_nat(other)),
384        Equal.trans(Maybe<&2, F32>, Core.read_value(input, Core.read_flat(input, Core.TTensor{kept_source(input, 0.0, history), source_fits}, other)),
385          Some{Replayed.value(input, 0.0, Replayed.snoc(history, StorageLaws.First{}), other)}, Some{Replayed.value(input, 0.0, history, other)},
386          Replayed.read_found(input, 0.0, Replayed.snoc(history, StorageLaws.First{}), source_fits, other, inside),
387          Equal.cong(F32, Maybe<&2, F32>, word => Some{word}, Replayed.value(input, 0.0, Replayed.snoc(history, StorageLaws.First{}), other),
388            Replayed.value(input, 0.0, history, other), half_keeps(input, 0.0, history, other))),
389        IndexMap.stored_at(input, values, counted, other, inside)), fits)
390
391def permute_contract(+input: Shape.Shape, +values: List<&2, F32>,
392  counted: {List.length(&2, F32, values) == Shape.logical_numel(Shape.dimensions(input)) : Nat},
393  +axes: List<&2, U32>, valid: {True{} == Layout.bijective(input, axes) : Bool}) -> Records.PermuteContract<input, values, axes>:
394  +counted = counted
395  +valid = valid
396  Records.PermuteContract{
397    FromValues.from_values_lift(input, values, counted,
398      result => {Observation.failure_code(input, Observation.Permute{axes}, result) == Some{0} : Maybe<&2, U32>},
399      fits => permuted_code(input, values, axes, valid, fits)),
400    FromValues.from_values_lift(input, values, counted,
401      result => {Observation.applied_dimensions(input, Observation.Permute{axes}, result) ==
402        Some{Axis.order(U32, Shape.dimensions(input), Shape.naturals(axes))} : Maybe<&2, List<&2, U32>>},
403      fits => permuted_dimensions(input, values, axes, valid, fits)),
404    FromValues.from_values_lift(input, values, counted,
405      result => {Observation.applied_length(input, Observation.Permute{axes}, result) ==
406        Some{Shape.logical_numel(Shape.dimensions(input))} : Maybe<&2, Nat>},
407      fits => permuted_length(input, values, axes, valid, fits)),
408    FromValues.from_values_lift(input, values, counted,
409      result => {Observation.applied_count(input, Observation.Permute{axes}, result) == Some{Shape.count(input)} : Maybe<&2, U32>},
410      fits => permuted_count(input, values, axes, valid, fits)),
411    index => inside =>
412      +inside = inside
413      FromValues.from_values_lift(input, values, counted,
414        result => {Observation.applied_value(input, Observation.Permute{axes}, index, result) ==
415          Values.listed(values, Values.permute_source(input, axes, index)) : Maybe<&2, F32>},
416        fits => permuted_value(input, values, counted, axes, valid, index, inside, fits)),
417    other => inside =>
418      +inside = inside
419      FromValues.from_values_lift(input, values, counted,
420        result => {Observation.kept_value(input, Observation.Permute{axes}, other, result) ==
421          List.get(&2, F32, values, U32.to_nat(other)) : Maybe<&2, F32>},
422        fits => permuted_kept(input, values, counted, axes, valid, other, inside, fits))}
423
424def requested(+input: Shape.Shape, +a: U32, +b: U32) -> List<&2, U32>:
425  Transpose.swapped_axes(input, a, b)
426
427def transpose_route(+input: Shape.Shape, +a: U32, +b: U32, source: Core.Tensor<input>,
428  decision: Shape.BoolCase<Transpose.valid(input, a, b)>, +inside: {True{} == Transpose.valid(input, a, b) : Bool}) ->
429  {Transforms.permute_copy(input, source, Layout.transpose_axes(Shape.dimensions(input), a, b, 0)) == Transforms.transpose_decided(input, source, a, b, decision) :
430    Core.Tensor<input> & Result<&2, &1, Transforms.Error, Core.Dynamic>}:
431  match decision:
432    case Shape.BFalse{equal}:
433      Empty.absurd({Transforms.permute_copy(input, source, Layout.transpose_axes(Shape.dimensions(input), a, b, 0)) ==
434        (source, Fail{Transforms.LayoutError{Layout.AxisOutOfBounds{}}}) : Core.Tensor<input> & Result<&2, &1, Transforms.Error, Core.Dynamic>},
435        Arith.clash(Equal.trans(Bool, False{}, Transpose.valid(input, a, b), True{}, equal, Equal.sym(Bool, True{}, Transpose.valid(input, a, b), inside))))
436    case Shape.BTrue{equal}:
437      {==}
438
439def requested_valid(+input: Shape.Shape, +a: U32, +b: U32, +inside: {True{} == Transpose.valid(input, a, b) : Bool}) ->
440  {True{} == Layout.bijective(input, requested(input, a, b)) : Bool}:
441  %TransposeProof.transpose_axes_requested(input, a, b, inside) : {True{} == Layout.bijective(input, _) : Bool}
442  TransposeProof.valid_axes(input, a, b, inside)
443
444# transpose_copy of any source with both axes in range is permute_copy by
445# the swapped axes.
446def transpose_cast(+input: Shape.Shape, +a: U32, +b: U32, +inside: {True{} == Transpose.valid(input, a, b) : Bool},
447  source: Core.Tensor<input>, -Q: Core.Tensor<input> & Result<&2, &1, Transforms.Error, Core.Dynamic> -> Type,
448  h: Q(Transforms.permute_copy(input, source, requested(input, a, b)))) -> Q(Transforms.transpose_copy(input, source, a, b)):
449  %transpose_route(input, a, b, source, Shape.bool_case(Transpose.valid(input, a, b)), inside) : Q(_)
450  %Equal.sym(List<&2, U32>, Layout.transpose_axes(Shape.dimensions(input), a, b, 0), requested(input, a, b),
451    TransposeProof.transpose_axes_requested(input, a, b, inside)) :
452    Q(Transforms.permute_copy(input, source, _))
453  h
454
455def transpose_lift(+input: Shape.Shape, +values: List<&2, F32>, +a: U32, +b: U32, +inside: {True{} == Transpose.valid(input, a, b) : Bool},
456  -fits: {Lifecycle.logical_length(Replayed.storage(input, 0.0, FromValues.writes(values, 0n))) == Shape.logical_numel(Shape.dimensions(input)) : Nat},
457  -Q: Core.Tensor<input> & Result<&2, &1, Transforms.Error, Core.Dynamic> -> Type,
458  h: Permuted(input, values, requested(input, a, b), fits, Q)) ->
459  Q(Transforms.transpose_copy(input, Core.TTensor{Replayed.storage(input, 0.0, FromValues.writes(values, 0n)), fits}, a, b)):
460  transpose_cast(input, a, b, inside, Core.TTensor{Replayed.storage(input, 0.0, FromValues.writes(values, 0n)), fits}, Q, h)
461
462def transpose_contract(+input: Shape.Shape, +values: List<&2, F32>,
463  counted: {List.length(&2, F32, values) == Shape.logical_numel(Shape.dimensions(input)) : Nat},
464  +a: U32, +b: U32, inside: {True{} == Transpose.valid(input, a, b) : Bool}) -> Records.TransposeContract<input, values, a, b>:
465  +counted = counted
466  +inside = inside
467  +axes = requested(input, a, b)
468  +valid = requested_valid(input, a, b, inside)
469  Records.TransposeContract{
470    FromValues.from_values_lift(input, values, counted,
471      result => {Observation.failure_code(input, Observation.Transpose{a, b}, result) == Some{0} : Maybe<&2, U32>},
472      fits => transpose_lift(input, values, a, b, inside, fits, pair => {Some{Observation.failure_code_pair(input, pair)} == Some{0} : Maybe<&2, U32>},
473        permuted_code(input, values, axes, valid, fits))),
474    FromValues.from_values_lift(input, values, counted,
475      result => {Observation.applied_dimensions(input, Observation.Transpose{a, b}, result) ==
476        Some{Axis.order(U32, Shape.dimensions(input), Shape.naturals(axes))} : Maybe<&2, List<&2, U32>>},
477      fits => transpose_lift(input, values, a, b, inside, fits, pair => {Observation.applied_dimensions_pair(input, pair) ==
478          Some{Axis.order(U32, Shape.dimensions(input), Shape.naturals(axes))} : Maybe<&2, List<&2, U32>>},
479        permuted_dimensions(input, values, axes, valid, fits))),
480    FromValues.from_values_lift(input, values, counted,
481      result => {Observation.applied_length(input, Observation.Transpose{a, b}, result) ==
482        Some{Shape.logical_numel(Shape.dimensions(input))} : Maybe<&2, Nat>},
483      fits => transpose_lift(input, values, a, b, inside, fits, pair => {Observation.applied_length_pair(input, pair) ==
484          Some{Shape.logical_numel(Shape.dimensions(input))} : Maybe<&2, Nat>},
485        permuted_length(input, values, axes, valid, fits))),
486    FromValues.from_values_lift(input, values, counted,
487      result => {Observation.applied_count(input, Observation.Transpose{a, b}, result) == Some{Shape.count(input)} : Maybe<&2, U32>},
488      fits => transpose_lift(input, values, a, b, inside, fits, pair => {Observation.applied_count_pair(input, pair) == Some{Shape.count(input)} : Maybe<&2, U32>},
489        permuted_count(input, values, axes, valid, fits))),
490    index => below =>
491      +below = below
492      FromValues.from_values_lift(input, values, counted,
493        result => {Observation.applied_value(input, Observation.Transpose{a, b}, index, result) ==
494          Values.listed(values, Values.permute_source(input, axes, index)) : Maybe<&2, F32>},
495        fits => transpose_lift(input, values, a, b, inside, fits, pair => {Observation.applied_value_pair(input, index, pair) ==
496            Values.listed(values, Values.permute_source(input, axes, index)) : Maybe<&2, F32>},
497          permuted_value(input, values, counted, axes, valid, index, below, fits))),
498    other => below =>
499      +below = below
500      FromValues.from_values_lift(input, values, counted,
501        result => {Observation.kept_value(input, Observation.Transpose{a, b}, other, result) ==
502          List.get(&2, F32, values, U32.to_nat(other)) : Maybe<&2, F32>},
503        fits => transpose_lift(input, values, a, b, inside, fits, pair => {Observation.kept_value_pair(input, other, pair) ==
504            List.get(&2, F32, values, U32.to_nat(other)) : Maybe<&2, F32>},
505          permuted_kept(input, values, counted, axes, valid, other, below, fits)))}
506
507# The words the copy loop writes, in index order, are `permuted_words`.
508def words_moved(+input: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>, +axes: List<&2, U32>, +plan: Layout.Plan<input, axes>,
509  +count: Nat, +room: {Nat.is_le(count, Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan)))) == True{} : Bool}) ->
510  {Contracts.words(Layout.output(input, axes, plan), 0.0,
511    IndexMap.writes(~Request, ~source_of, Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan))), Request{input, axes}, input, fill, ops),
512    count) == Records.permuted_words(input, fill, ops, axes, count) : List<&2, F32>}:
513  match count:
514    case 0n:
515      {==}
516    case 1n+previous:
517      +output = Layout.output(input, axes, plan)
518      +total = Shape.logical_numel(Shape.dimensions(output))
519      +request : Request = Request{input, axes}
520      +history = IndexMap.writes(~Request, ~source_of, total, request, input, fill, ops)
521      +below = Conversion.le_successor_lt(previous, total, room)
522      +position = U32.from_nat(previous)
523      Words.snoc_equal(Contracts.words(output, 0.0, history, previous), Records.permuted_words(input, fill, ops, axes, previous),
524        Replayed.value(output, 0.0, history, position),
525        Replayed.value(input, fill, ops, U32.from_nat(Values.permute_source_index(input, axes, previous))),
526        words_moved(input, fill, ops, axes, plan, previous, Conversion.le_predecessor(previous, total, room)),
527        Equal.trans(F32, Replayed.value(output, 0.0, history, position),
528          Replayed.value(input, fill, ops, U32.from_nat(Values.permute_source_index(input, axes, U32.to_nat(position)))),
529          Replayed.value(input, fill, ops, U32.from_nat(Values.permute_source_index(input, axes, previous))),
530          PointwiseProof.writes_value(~Unit, ~IndexMap.keep, ~Request, ~source_of, ~source_of, total, Unit{}, request, input, fill, ops,
531            input, fill, ops, output, position, PointwiseProof.word_inside(previous, output, below), Conversion.le_reflexive(total)),
532          Equal.cong(Nat, F32, number => Replayed.value(input, fill, ops, U32.from_nat(Values.permute_source_index(input, axes, number))),
533            U32.to_nat(position), previous, PointwiseProof.word_exact(previous, output, below))))
534
535# The copy loop on any replayed source, in transport form: once
536# `Layout.permutation` gives the plan, P holds for the permute if it holds for
537# every result of the form the loop's final state has.
538def replayed_final(+input: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>, +axes: List<&2, U32>, +plan: Layout.Plan<input, axes>,
539  -P: Core.Dynamic & Result<&2, &1, Transforms.Error, Core.Dynamic> -> Type,
540  h: Records.PermutedWith(input, fill, ops, axes, plan, P)) ->
541  Final(input, axes, plan, fill, ops, Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan))),
542    pair => P(Transforms.dynamic_result(input, Transforms.copy_pack(input, axes, plan, pair)))):
543  +output = Layout.output(input, axes, plan)
544  +total = Shape.logical_numel(Shape.dimensions(output))
545  +request : Request = Request{input, axes}
546  source_fits => target_fits =>
547    h(Replayed.snoc(ops, StorageLaws.First{}), words_first(input, fill, ops, Shape.logical_numel(Shape.dimensions(input))),
548      source_fits,
549      IndexMap.writes(~Request, ~source_of, total, request, input, fill, ops),
550      words_moved(input, fill, ops, axes, plan, total, Conversion.le_reflexive(total)),
551      target_fits)
552
553def permute_replayed(+input: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>,
554  -fits: {Lifecycle.logical_length(Replayed.storage(input, fill, ops)) == Shape.logical_numel(Shape.dimensions(input)) : Nat},
555  +axes: List<&2, U32>, +plan: Layout.Plan<input, axes>,
556  laid: {Layout.permutation(input, axes) == Done{plan} : Result<&2, &2, Layout.Error, Layout.Plan<input, axes>>},
557  -P: Core.Dynamic & Result<&2, &1, Transforms.Error, Core.Dynamic> -> Type,
558  h: Records.PermutedWith(input, fill, ops, axes, plan, P)) ->
559  P(Transforms.dynamic_result(input, Transforms.permute_copy(input, Core.TTensor{Replayed.storage(input, fill, ops), fits}, axes))):
560  %Equal.sym(Result<&2, &2, Layout.Error, Layout.Plan<input, axes>>, Layout.permutation(input, axes), Done{plan}, laid) :
561    P(Transforms.dynamic_result(input, Transforms.copy_planned(input, axes, Core.TTensor{Replayed.storage(input, fill, ops), fits}, _)))
562  run_lift(input, axes, plan, fill, ops, pair => P(Transforms.dynamic_result(input, pair)),
563    replayed_final(input, fill, ops, axes, plan, P, h), fits)
564
565def transpose_replayed(+input: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>,
566  -fits: {Lifecycle.logical_length(Replayed.storage(input, fill, ops)) == Shape.logical_numel(Shape.dimensions(input)) : Nat},
567  +a: U32, +b: U32, inside: {True{} == Transpose.valid(input, a, b) : Bool},
568  +plan: Layout.Plan<input, Transpose.swapped_axes(input, a, b)>,
569  laid: {Layout.permutation(input, Transpose.swapped_axes(input, a, b)) == Done{plan} :
570    Result<&2, &2, Layout.Error, Layout.Plan<input, Transpose.swapped_axes(input, a, b)>>},
571  -P: Core.Dynamic & Result<&2, &1, Transforms.Error, Core.Dynamic> -> Type,
572  h: Records.PermutedWith(input, fill, ops, Transpose.swapped_axes(input, a, b), plan, P)) ->
573  P(Transforms.dynamic_result(input, Transforms.transpose_copy(input, Core.TTensor{Replayed.storage(input, fill, ops), fits}, a, b))):
574  +inside = inside
575  transpose_cast(input, a, b, inside, Core.TTensor{Replayed.storage(input, fill, ops), fits}, pair => P(Transforms.dynamic_result(input, pair)),
576    permute_replayed(input, fill, ops, fits, requested(input, a, b), plan, laid, P, h))