1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
/-! # Permutations as data

A permutation of the entries of a vector of length `n` is a sequence of steps (`PermS`), each a
numpy operation on the vector reshaped, row-major, to a shape `dims` with `n` entries:

| step | numpy |
|---|---|
| `tr dims axes` | `x.reshape(dims).transpose(axes).reshape(-1)` |
| `flip dims k` | `jnp.flip(x.reshape(dims), k).reshape(-1)` |
| `roll dims k s` | `jnp.roll(x.reshape(dims), s, axis=k).reshape(-1)` |

The steps apply in order. Each acts as a gather: output position `i` reads input position
`idx i`, so the sequence `[p₁, p₂, …]` reads `p₁.idx (p₂.idx (… i))` (`PermS.idx`). The squeezes
and patchings of the models are transposes, their channel and token reversals flips, and their
rotations rolls.

Nothing here uses Mathlib: the flow signature (`Flow.Basic`) reads `PermS.idx` as an
`Equiv.Perm`, and the printer (`Flow.Print`) prints the steps. -/

namespace RoleFlow.Flow

/-- One step of a permutation, on the vector reshaped to `dims`. -/
inductive PermStep
  /-- Permute the axes: output axis `j` is input axis `axes[j]`. -/
  | tr (dims axes : List Nat)
  /-- Reverse axis `k`. -/
  | flip (dims : List Nat) (k : Nat)
  /-- Rotate axis `k` by `s`: output index `j` reads `j - s` (mod the axis length). -/
  | roll (dims : List Nat) (k s : Nat)
  deriving DecidableEq, Repr, Inhabited

/-- A permutation: its steps, applied in order. -/
abbrev PermS := List PermStep

namespace PermStep

/-- The row-major digits of `i` in the shape `dims`. -/
def digits (dims : List Nat) (i : Nat) : List Nat :=
  (dims.foldr (fun d (acc, j) => (j % d :: acc, j / d)) ([], i)).1

/-- The position of the digits `ds` in the shape `dims`. -/
def undigits (dims ds : List Nat) : Nat :=
  (dims.zip ds).foldl (fun acc (d, x) => acc * d + x) 0

def dims : PermStep → List Nat
  | .tr d _ | .flip d _ | .roll d _ _ => d

/-- The input position that output position `i` reads. -/
def idx : PermStep → Nat → Nat
  | .tr dims axes, i =>
    let e := digits (axes.map fun a => dims.getD a 1) i
    undigits dims ((List.range dims.length).map fun a => e.getD (axes.idxOf a) 0)
  | .flip dims k, i =>
    let ds := digits dims i
    undigits dims (ds.set k (dims.getD k 1 - 1 - ds.getD k 0))
  | .roll dims k s, i =>
    let ds := digits dims i
    let d := dims.getD k 1
    undigits dims (ds.set k ((ds.getD k 0 + d - s % d) % d))

/-- The step that undoes this one. -/
def inv : PermStep → PermStep
  | .tr dims axes => .tr (axes.map fun a => dims.getD a 1) ((List.range axes.length).map axes.idxOf)
  | .flip dims k => .flip dims k
  | .roll dims k s => let d := dims.getD k 1; .roll dims k ((d - s % d) % d)

/-- The step is well formed on vectors of length `n`: its shape has `n` entries, `axes` is a
permutation of the axes, and `k` is an axis. -/
def ok (n : Nat) (p : PermStep) : Bool :=
  p.dims.foldl (· * ·) 1 == n && match p with
    | .tr dims axes => axes.length == dims.length && (List.range dims.length).all axes.contains
    | .flip dims k | .roll dims k _ => k < dims.length

/-- The step on the middle of the shape `[outer] ++ dims ++ [inner]` (an axis of length `1` is
left out). -/
def embed (outer inner : Nat) (p : PermStep) : PermStep :=
  let pre := if outer = 1 then [] else [outer]
  let post := if inner = 1 then [] else [inner]
  let o := pre.length
  match p with
  | .tr dims axes =>
    .tr (pre ++ dims ++ post) ((List.range o) ++ axes.map (· + o) ++ (post.map fun _ => dims.length + o))
  | .flip dims k => .flip (pre ++ dims ++ post) (k + o)
  | .roll dims k s => .roll (pre ++ dims ++ post) (k + o) s

/-- The step does nothing: a transpose with the axes in order, a flip of an axis of length at most
`1`, a roll by a multiple of the axis length. -/
def isId : PermStep → Bool
  | .tr _ axes => axes == List.range axes.length
  | .flip dims k => dims.getD k 1 ≤ 1
  | .roll dims k s => s % dims.getD k 1 == 0

/-- The step in the surface syntax of the webapp, before its argument: `tr<[2, 3], [1, 0]>`. -/
def text : PermStep → String
  | .tr dims axes => s!"tr<{dims}, {axes}>"
  | .flip dims k => s!"flip<{dims}, {k}>"
  | .roll dims k s => s!"roll<{dims}, {k}, {s}>"

end PermStep

namespace PermS

/-- The steps applied to `x`, in the surface syntax of the webapp. -/
def text (x : String) (ps : PermS) : String := ps.foldl (fun x p => s!"{p.text}({x})") x

/-- The input position that output position `i` reads. -/
def idx (ps : PermS) (i : Nat) : Nat := ps.foldr PermStep.idx i

/-- The steps that undo `ps`. -/
def inv (ps : PermS) : PermS := ps.reverse.map PermStep.inv

/-- Every step is well formed on vectors of length `n`. -/
def ok (n : Nat) (ps : PermS) : Bool := ps.all (PermStep.ok n)

/-- The index list `[idx 0, …, idx (n - 1)]`. -/
def idxList (n : Nat) (ps : PermS) : List Nat := (List.range n).map ps.idx

/-- `ps` is a permutation of `0, …, n - 1`: well formed, and its index list is one. -/
def valid (n : Nat) (ps : PermS) : Bool :=
  ps.ok n && (let l := ps.idxList n
    let seen := l.foldl (fun (a : Array Bool) j => if j < n then a.set! j true else a)
      (Array.replicate n false)
    l.all (· < n) && seen.all id)

end PermS

end RoleFlow.Flow