def flatten_unflatten

bendtorch-core@0.1.0.0 · src/shape/index_proofs.bend:123

import Base
import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape

123def flatten_unflatten(+dims: List<&2, Nat>, +index: Shape.Fin(Shape.numel(dims))) -> {Shape.flatten(dims, Shape.unflatten(dims, index)) == index : Shape.Fin(Shape.numel(dims))}:
124  match dims:
125    case Nil{}:
126      match index:
127        case Shape.FZero{}:
128          {==}
129        case Shape.FNext{previous}:
130          match previous:
131    case +head <> +tail:
132      %pack_unpack(head, Shape.numel(tail), index) : {Shape.flatten(head <> tail, Shape.unflatten(head <> tail, index)) == _ : Shape.Fin(Shape.numel(head <> tail))}
133      flatten_unflatten_step(head, tail, +last => flatten_unflatten(tail, last), Shape.Fin.unpack(head, Shape.numel(tail), index))