src/tensor/permute_values_records.bend
bendtorch-core@0.1.0.0 · 72 lines · 4 definitions
1import Base
2import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
3import bendtorch-core@0.1.0.0/src/tensor/core.bend as Core
4import bendtorch-core@0.1.0.0/src/tensor/axis_spec.bend as Axis
5import bendtorch-core@0.1.0.0/src/tensor/transpose_spec.bend as Transpose
6import bendtorch-core@0.1.0.0/src/tensor/observation.bend as Values
7import bendtorch-core@0.1.0.0/src/tensor/transform_observation.bend as Observation
8import bendtorch-core@0.1.0.0/src/storage/LAWS.bend as StorageLaws
9import bendtorch-core@0.1.0.0/src/tensor/lifecycle.bend as Lifecycle
10import bendtorch-core@0.1.0.0/src/tensor/layout.bend as Layout
11import bendtorch-core@0.1.0.0/src/tensor/transforms.bend as Transforms
12import bendtorch-core@0.1.0.0/src/tensor/contracts/replayed.bend as Replayed
13import bendtorch-core@0.1.0.0/src/tensor/contracts/LAWS.bend as Contracts
14
15type PermuteContract<-input: Shape.Shape, -values: List<&2, F32>, -axes: List<&2, U32>> is Type:
16 PermuteContract{
17 outcome: {Observation.failure_code(input, Observation.Permute{axes}, Core.from_values(input, values)) == Some{0} : Maybe<&2, U32>},
18 dimensions: {Observation.applied_dimensions(input, Observation.Permute{axes}, Core.from_values(input, values)) ==
19 Some{Axis.order(U32, Shape.dimensions(input), Shape.naturals(axes))} : Maybe<&2, List<&2, U32>>},
20 count: {Observation.applied_length(input, Observation.Permute{axes}, Core.from_values(input, values)) ==
21 Some{Shape.logical_numel(Shape.dimensions(input))} : Maybe<&2, Nat>},
22 word_count: {Observation.applied_count(input, Observation.Permute{axes}, Core.from_values(input, values)) == Some{Shape.count(input)} : Maybe<&2, U32>},
23 value: @+index: U32 -> @inside: {Nat.is_lt(U32.to_nat(index), Shape.logical_numel(Shape.dimensions(input))) == True{} : Bool} ->
24 {Observation.applied_value(input, Observation.Permute{axes}, index, Core.from_values(input, values)) ==
25 Values.listed(values, Values.permute_source(input, axes, index)) : Maybe<&2, F32>},
26 retained_value: @+other: U32 -> @inside: {Nat.is_lt(U32.to_nat(other), Shape.logical_numel(Shape.dimensions(input))) == True{} : Bool} ->
27 {Observation.kept_value(input, Observation.Permute{axes}, other, Core.from_values(input, values)) ==
28 List.get(&2, F32, values, U32.to_nat(other)) : Maybe<&2, F32>}}
29
30type TransposeContract<-input: Shape.Shape, -values: List<&2, F32>, -a: U32, -b: U32> is Type:
31 TransposeContract{
32 outcome: {Observation.failure_code(input, Observation.Transpose{a, b}, Core.from_values(input, values)) == Some{0} : Maybe<&2, U32>},
33 dimensions: {Observation.applied_dimensions(input, Observation.Transpose{a, b}, Core.from_values(input, values)) ==
34 Some{Axis.order(U32, Shape.dimensions(input), Shape.naturals(Transpose.requested_axes(Shape.rank_dimensions(Shape.dimensions(input)), U32.to_nat(a), U32.to_nat(b))))} : Maybe<&2, List<&2, U32>>},
35 count: {Observation.applied_length(input, Observation.Transpose{a, b}, Core.from_values(input, values)) ==
36 Some{Shape.logical_numel(Shape.dimensions(input))} : Maybe<&2, Nat>},
37 word_count: {Observation.applied_count(input, Observation.Transpose{a, b}, Core.from_values(input, values)) == Some{Shape.count(input)} : Maybe<&2, U32>},
38 value: @+index: U32 -> @inside: {Nat.is_lt(U32.to_nat(index), Shape.logical_numel(Shape.dimensions(input))) == True{} : Bool} ->
39 {Observation.applied_value(input, Observation.Transpose{a, b}, index, Core.from_values(input, values)) ==
40 Values.listed(values, Values.permute_source(input, Transpose.requested_axes(Shape.rank_dimensions(Shape.dimensions(input)), U32.to_nat(a), U32.to_nat(b)), index)) : Maybe<&2, F32>},
41 retained_value: @+other: U32 -> @inside: {Nat.is_lt(U32.to_nat(other), Shape.logical_numel(Shape.dimensions(input))) == True{} : Bool} ->
42 {Observation.kept_value(input, Observation.Transpose{a, b}, other, Core.from_values(input, values)) ==
43 List.get(&2, F32, values, U32.to_nat(other)) : Maybe<&2, F32>}}
44
45
46
47def permuted_words(+input: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>, +axes: List<&2, U32>, +count: Nat) -> List<&2, F32>:
48 match count:
49 case 0n:
50 Nil{}
51 case 1n+previous:
52 List.append(&2, F32, permuted_words(input, fill, ops, axes, previous),
53 [Replayed.value(input, fill, ops, U32.from_nat(Values.permute_source_index(input, axes, previous)))])
54
55
56
57
58
59def PermutedWith(+input: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>, +axes: List<&2, U32>, +plan: Layout.Plan<input, axes>,
60 -P: Core.Dynamic & Result<&2, &1, Transforms.Error, Core.Dynamic> -> Type) -> Type:
61 @+kept: List<&2, StorageLaws.Op> ->
62 @kept_words: {Contracts.words(input, fill, kept, Shape.logical_numel(Shape.dimensions(input))) ==
63 Contracts.words(input, fill, ops, Shape.logical_numel(Shape.dimensions(input))) : List<&2, F32>} ->
64 @-kept_fits: {Lifecycle.logical_length(Replayed.storage(input, fill, kept)) == Shape.logical_numel(Shape.dimensions(input)) : Nat} ->
65 @+output_ops: List<&2, StorageLaws.Op> ->
66 @moved: {Contracts.words(Layout.output(input, axes, plan), 0.0, output_ops, Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan)))) ==
67 permuted_words(input, fill, ops, axes, Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan)))) : List<&2, F32>} ->
68 @-output_fits: {Lifecycle.logical_length(Replayed.storage(Layout.output(input, axes, plan), 0.0, output_ops)) ==
69 Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan))) : Nat} ->
70 P((Core.Packed{input, Core.TTensor{Replayed.storage(input, fill, kept), kept_fits}},
71 Done{Core.Packed{Layout.output(input, axes, plan),
72 Core.TTensor{Replayed.storage(Layout.output(input, axes, plan), 0.0, output_ops), output_fits}}}))