def unflatten_at

bendtorch-core@0.1.0.0 · src/tensor/matrix_index_proof.bend:132

import Base
import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
import bendtorch-core@0.1.0.0/src/shape/LAWS.bend as ShapeLaws
import bendtorch-core@0.1.0.0/src/shape/boundary_proofs.bend as Boundary
import bendtorch-core@0.1.0.0/src/shape/coordinate_proofs.bend as Coordinates
import bendtorch-core@0.1.0.0/src/tensor/matmul.bend as Matrix

132def unflatten_at(+output: Shape.Shape, +rows: Nat, +columns: Nat,
133  extents: {Shape.extents(output) == [rows, columns] : List<&2, Nat>}, +position: Nat,
134  +below: {Nat.is_lt(position, Shape.logical_numel(Shape.dimensions(output))) == True{} : Bool}) ->
135  {ShapeLaws.coordinate_result(output, Shape.unflatten_checked(output, U32.from_nat(position))) ==
136    Some{[U32.from_nat(Matrix.row_of(columns, position)), U32.from_nat(Matrix.column_of(columns, position))]} : Maybe<&2, List<&2, U32>>}:
137  +bound = Shape.numel(Shape.extents(output))
138  +fits = Equal.sym(Bool, Nat.is_lt(position, bound), True{}, below)
139  +index = Coordinates.fin_from(bound, position, fits)
140  Equal.trans(Maybe<&2, List<&2, U32>>, ShapeLaws.coordinate_result(output, Shape.unflatten_checked(output, U32.from_nat(position))),
141    ShapeLaws.unflatten_spec(Shape.extents(output), ShapeLaws.fin_of(bound, U32.to_nat(U32.from_nat(position)))),
142    Some{[U32.from_nat(Matrix.row_of(columns, position)), U32.from_nat(Matrix.column_of(columns, position))]},
143    Boundary.unflatten_checked(output, U32.from_nat(position)),
144    Equal.trans(Maybe<&2, List<&2, U32>>,
145      ShapeLaws.unflatten_spec(Shape.extents(output), ShapeLaws.fin_of(bound, U32.to_nat(U32.from_nat(position)))),
146      ShapeLaws.unflatten_spec(Shape.extents(output), ShapeLaws.fin_of(bound, position)),
147      Some{[U32.from_nat(Matrix.row_of(columns, position)), U32.from_nat(Matrix.column_of(columns, position))]},
148      Equal.cong(Nat, Maybe<&2, List<&2, U32>>, value => ShapeLaws.unflatten_spec(Shape.extents(output), ShapeLaws.fin_of(bound, value)),
149        U32.to_nat(U32.from_nat(position)), position, word_exact(position, output, below)),
150      Equal.trans(Maybe<&2, List<&2, U32>>,
151        ShapeLaws.unflatten_spec(Shape.extents(output), ShapeLaws.fin_of(bound, position)),
152        ShapeLaws.unflatten_spec(Shape.extents(output), Some{index}),
153        Some{[U32.from_nat(Matrix.row_of(columns, position)), U32.from_nat(Matrix.column_of(columns, position))]},
154        Equal.cong(Maybe<&2, Shape.Fin(bound)>, Maybe<&2, List<&2, U32>>, found => ShapeLaws.unflatten_spec(Shape.extents(output), found),
155          ShapeLaws.fin_of(bound, position), Some{index}, Boundary.fin_of_from(bound, position, fits)),
156        Equal.cong(List<&2, Nat>, Maybe<&2, List<&2, U32>>, values => Some{Shape.words(values)},
157          Shape.indices(Shape.extents(output), Shape.unflatten(Shape.extents(output), index)),
158          [Matrix.row_of(columns, position), Matrix.column_of(columns, position)],
159          position_digits(output, rows, columns, extents, position, fits)))))