def position_fits

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

import Base
import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
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

161def position_fits(+output: Shape.Shape, +rows: Nat, +columns: Nat,
162  +extents: {Shape.extents(output) == [rows, columns] : List<&2, Nat>}, +position: Nat,
163  +below: {Nat.is_lt(position, Shape.logical_numel(Shape.dimensions(output))) == True{} : Bool}) ->
164  {Shape.fit([rows, columns], [Matrix.row_of(columns, position), Matrix.column_of(columns, position)]) == True{} : Bool}:
165  +bound = Shape.numel(Shape.extents(output))
166  +fits = Equal.sym(Bool, Nat.is_lt(position, bound), True{}, below)
167  +coordinate = Shape.unflatten(Shape.extents(output), Coordinates.fin_from(bound, position, fits))
168  +digits : List<&2, Nat> = [Matrix.row_of(columns, position), Matrix.column_of(columns, position)]
169  Equal.sym(Bool, True{}, Shape.fit([rows, columns], digits),
170    Equal.trans(Bool, True{}, Shape.fit(Shape.extents(output), Shape.indices(Shape.extents(output), coordinate)), Shape.fit([rows, columns], digits),
171      Coordinates.fit_indices(Shape.extents(output), coordinate),
172      Equal.trans(Bool, Shape.fit(Shape.extents(output), Shape.indices(Shape.extents(output), coordinate)),
173        Shape.fit(Shape.extents(output), digits), Shape.fit([rows, columns], digits),
174        Equal.cong(List<&2, Nat>, Bool, values => Shape.fit(Shape.extents(output), values),
175          Shape.indices(Shape.extents(output), coordinate), digits, position_digits(output, rows, columns, extents, position, fits)),
176        Equal.cong(List<&2, Nat>, Bool, dims => Shape.fit(dims, digits), Shape.extents(output), [rows, columns], extents))))