The JAX printer (trusted)

A typed program P of the flow signature, with its roles ρ, is printed as two JAX functions that also return the log-weight: the forward function from the lines of P, and the inverse from the lines of the inverse program P† (Prog.inverse, the schedule of the recipes of the derivation). The inverse program is typed too, with its own roles ρ† (inferred), and both are printed by the same code (emitFwd): the log-weight is L_θ, the sum of ℓ_o over the spine steps, accumulated in ld with one ld += ℓ_o on each spine step's row, the ; aligned. For P† it is the log-weight of Prog.density for P†; that it agrees a.e. with -L_θ at the values P† recomputes (Prog.inverse_density) is not proved (TODO.md), but checked by scripts/compare.py. Every variable is a Python expression; structural operations (split, the changes of sort) emit slices or nothing, the others bind a fresh variable; a spine step that prints nothing gets a row pass ; ld += ….

Each stage is a definition (Linked.add) printed as {name}_fwd, {name}_inv over its own weights. A model is printed as one program flow_fwd (and its inverse program flow_inv), so the stage pipeline is a typed program rather than a Python composition: its stages inlined (Flow.modelProg), or, with calls, a call of each stage (Flow.callProg), a row y, l = {name}_fwd(w[a:b], x) ; ld += l. A call is a spine step like any other, and the inverse program calls {name}_inv. The roles are inferred (Prog.infer) on each stage and on the model. writeModule refuses a model or a stage that is not typed, or whose inverse program is not.

The leaf table (one JAX expression per operation and per nonzero log-weight) is trusted: it is written to agree with Flow.semOp and Flow.ℓOp. Functions act on one example; the caller batches with jax.vmap.

The printer's state

structure RoleFlow.Flow.Print.Gen : Type

• lines : Array (String × String)

  The rows so far: the code, and the log-weight statement after its ; ("": none).

• next : ℕ
• permChecks : Array (String × ℕ × List ℕ)

  Each permutation's expression (on x, of length n) with its index list, for <module>_perms.py.

@[reducible, inline]

abbrev RoleFlow.Flow.Print.M (α : Type) : Type

Equations

• RoleFlow.Flow.Print.M = StateM RoleFlow.Flow.Print.Gen

def RoleFlow.Flow.Print.emitLine (l : String) : M Unit

Equations

• RoleFlow.Flow.Print.emitLine l = modify fun (g : RoleFlow.Flow.Print.Gen) => { lines := g.lines.push (l, ""), next := g.next, permChecks := g.permChecks }

def RoleFlow.Flow.Print.bindE (e : String) : M String

Bind an expression to a fresh variable.

Equations

• One or more equations did not get rendered due to their size.

Permutations

A permutation is printed as its steps (Flow.PermS), one numpy call each, applied in order; its inverse as the steps of PermS.inv. writeModule also writes each permutation's expression with its index list (PermS.idxList) to <module>_perms.py, which checks that numpy agrees.

def RoleFlow.Flow.Print.numsE (l : List ℕ) : String

The comma-separated numbers.

Equations

• RoleFlow.Flow.Print.numsE l = ", ".intercalate (List.map toString l)

def RoleFlow.Flow.Print.stepE (x : String) : PermStep → String

A step applied to the expression x; the result keeps the step's shape (reshape reads row-major, so the next step can reshape it directly).

Equations

• One or more equations did not get rendered due to their size.

def RoleFlow.Flow.Print.permStepsE (x : String) (ps : PermS) : String

The steps applied to the flat vector x, in order, flattened at the end; steps that do nothing (PermStep.isId) are left out.

Equations

• RoleFlow.Flow.Print.permStepsE x ps = match List.filter (fun (x : RoleFlow.Flow.PermStep) => !x.isId) ps with | [] => x | qs => List.foldl RoleFlow.Flow.Print.stepE x qs ++ ".reshape(-1)"

def RoleFlow.Flow.Print.concatE : List String → String

The empty vector is printed as bool, which jnp.concatenate promotes to the dtype of what it is joined with.

Equations

• RoleFlow.Flow.Print.concatE [] = "jnp.zeros(0, bool)"
• RoleFlow.Flow.Print.concatE x✝ = toString "jnp.concatenate([" ++ toString (", ".intercalate x✝) ++ toString "])"

def RoleFlow.Flow.Print.tupE : List String → String

Several values as one Python value: a tuple, or the value itself.

Equations

• RoleFlow.Flow.Print.tupE [e] = e
• RoleFlow.Flow.Print.tupE x✝ = toString "(" ++ toString (", ".intercalate x✝) ++ toString ")"

Leaves

def RoleFlow.Flow.Print.logpdfE : Ref → Option String

The function computing Ref.logpdf, where it is not 0.

Equations

• RoleFlow.Flow.Print.logpdfE RoleFlow.Flow.Ref.leb = none
• RoleFlow.Flow.Print.logpdfE RoleFlow.Flow.Ref.normal = some "logphi"
• RoleFlow.Flow.Print.logpdfE RoleFlow.Flow.Ref.logistic = some "loglogistic"

def RoleFlow.Flow.Print.emitOp {as bs : List FS} : FOp as bs → List String → M (List String)

An operation's results from its arguments, one expression per result; agrees with Flow.semOp.

Equations

• One or more equations did not get rendered due to their size.
• RoleFlow.Flow.Print.emitOp RoleFlow.Flow.FOp.nil x✝ = do let __do_lift ← RoleFlow.Flow.Print.bindE "jnp.zeros(0, bool)" pure [__do_lift]
• RoleFlow.Flow.Print.emitOp (RoleFlow.Flow.FOp.concat m n) x✝ = do let __do_lift ← RoleFlow.Flow.Print.bindE (RoleFlow.Flow.Print.concatE x✝) pure [__do_lift]
• RoleFlow.Flow.Print.emitOp (RoleFlow.Flow.FOp.rebase r r' n) x✝ = pure x✝
• RoleFlow.Flow.Print.emitOp (RoleFlow.Flow.FOp.relu n) x✝ = do let __do_lift ← RoleFlow.Flow.Print.bindE (toString "jax.nn.relu(" ++ toString (x✝.headD "") ++ toString ")") pure [__do_lift]
• RoleFlow.Flow.Print.emitOp (RoleFlow.Flow.FOp.tanh n) x✝ = do let __do_lift ← RoleFlow.Flow.Print.bindE (toString "jnp.tanh(" ++ toString (x✝.headD "") ++ toString ")") pure [__do_lift]
• RoleFlow.Flow.Print.emitOp (RoleFlow.Flow.FOp.layernorm n) x✝ = do let __do_lift ← RoleFlow.Flow.Print.bindE (toString "layernorm(" ++ toString (x✝.headD "") ++ toString ")") pure [__do_lift]
• RoleFlow.Flow.Print.emitOp (RoleFlow.Flow.FOp.mul n) x✝ = do let __do_lift ← RoleFlow.Flow.Print.bindE (toString (x✝.getD 0 "") ++ toString " * " ++ toString (x✝.getD 1 "")) pure [__do_lift]

def RoleFlow.Flow.Print.ldOp {as bs : List FS} : FOp as bs → Option (List String → String)

The log-weight ℓ_o on the arguments, where it is not 0; agrees with Flow.ℓOp (rebase r r' has log ρ_{r'} - log ρ_r, a term for each density that is not 1; oddpow_logdet p x is Σ log p x^(p-1), which opowLogdet sets to 0 on the null set of the axes; triu_logdet c n U = n Σ_a U_aa is triuLogdet c n U).

Equations

• One or more equations did not get rendered due to their size.
• RoleFlow.Flow.Print.ldOp x✝ = none

def RoleFlow.Flow.Print.addLd (n : ℕ) (t : String) : M Unit

Put ld += t after the row the current line printed, if it printed one (the rows were n before it), or else on a row pass of its own.

Equations

• One or more equations did not get rendered due to their size.

Calls

def RoleFlow.Flow.Print.weightsE (env : Array String) (ps : List ℕ) : String

The weights at the positions ps, as a list: a slice w[a:b] if they are the function's own weights a, …, b - 1.

Equations

• One or more equations did not get rendered due to their size.
• RoleFlow.Flow.Print.weightsE env [] = "[]"

def RoleFlow.Flow.Print.emitCall (name : String) (inv : Bool) (ws x : String) (nb : ℕ) : M (List String × String)

A row calling {name}_fwd (or _inv) on the weights ws and the values x, binding its nb results and its log-weight; returns their expressions.

Equations

• One or more equations did not get rendered due to their size.

Programs

@[reducible, inline]

abbrev RoleFlow.Flow.Print.View (Sg : Sig FS) : Type

An operation of the signature the printed program is over: of the flow signature, or a call (Linked.view).

Equations

• RoleFlow.Flow.Print.View Sg = ({as bs : List RoleFlow.Flow.FS} → Sg.Op as bs → RoleFlow.Callee RoleFlow.Flow.flowSig.Op as bs)

def RoleFlow.Flow.Print.emitLines {Sg : Sig FS} (view : View Sg) (env : Array String) (ls : List (Line Sg)) (ρ : List Role) : M (Array String)

Print the lines from the expressions of the context they start in, each spine step with its ℓ_o (for a call, the log-weight it returns); returns the expressions of all variables, by position.

Equations

• One or more equations did not get rendered due to their size.

def RoleFlow.Flow.Print.ldCol : ℕ

The column of the ;: after the longest row that has a log-weight statement, ignoring rows longer than this.

Equations

• RoleFlow.Flow.Print.ldCol = 96

def RoleFlow.Flow.Print.emitDef (name : String) (body : M String) : M String

def name(w, x) returning (value, ld), from the body's rows: ld starts at 0, and the log-weight statements are aligned at one column.

Equations

• One or more equations did not get rendered due to their size.

def RoleFlow.Flow.Print.argsE (n : ℕ) : List String

The inputs of a function taking n values as x.

Equations

• RoleFlow.Flow.Print.argsE n = if n = 1 then ["x"] else List.map (fun (j : ℕ) => toString "x[" ++ toString j ++ toString "]") (List.range n)

def RoleFlow.Flow.Print.emitFwd {Sg : Sig FS} (view : View Sg) (P : Prog Sg) (ρ : List Role) (env₀ : List String) : M (List String)

The function of P from the expressions of its weights and inputs: the returned expressions, after the rows that compute them and accumulate the log-weight.

Equations

• One or more equations did not get rendered due to their size.

def RoleFlow.Flow.Print.emitPair {Sg : Sig FS} (view : View Sg) (name : String) (P : Prog Sg) (ρ ρi : List Role) : M String

{name}_fwd from P with its roles ρ, and {name}_inv from P† = P.inverse ρ with its roles ρi, both with the log-weight. The weights are the list w.

Equations

• One or more equations did not get rendered due to their size.

Modules

def RoleFlow.Flow.Print.preamble : String

Helpers the leaf table refers to, matching Flow.FlowOps.

Equations

• One or more equations did not get rendered due to their size.

def RoleFlow.Flow.Print.postamble (nOut : ℕ) : String

The postamble: the model's forward/inverse (the model as one program) and init_weights.

Equations

• One or more equations did not get rendered due to their size.

def RoleFlow.Flow.Print.emitStage {Sg : Sig FS} (view : View Sg) (s : Stage) (Q : Prog Sg) (ρ ρi : List Role) : M String

A stage's two functions (over its own weights), and its weights' dimensions and inits.

Equations

• One or more equations did not get rendered due to their size.

def RoleFlow.Flow.Print.emitDefs (stages : List Stage) {Sg : Sig FS} : Linked flowSig Sg → M (List String)

The definitions of a tower, in order, each with the dimensions and inits of the weights of the stage of its name.

Equations

• One or more equations did not get rendered due to their size.
• RoleFlow.Flow.Print.emitDefs stages RoleFlow.Linked.base = pure []

structure RoleFlow.Flow.Print.Cert : Type 1

A certified model: its stages, as the definitions of a tower L (each over its own weights), and the model as one program P at the top of the tower with its roles ρ and the roles ρi of P.inverse ρ: the stages inlined, or a program of calls of them.

• stages : List Stage
• Sg : Sig FS
• L : Linked flowSig self.Sg
• P : Prog self.Sg
• ρ : List Role
• ρi : List Role

def RoleFlow.Flow.Print.emitModuleChecks (c : Cert) : String × Array (String × ℕ × List ℕ)

The module: the stages, the model (flow_fwd from P, flow_inv from P†), and the Python interface; with each permutation's expression and index list.

Equations

• One or more equations did not get rendered due to their size.

def RoleFlow.Flow.Print.emitModule (c : Cert) : String

The module (emitModuleChecks).

Equations

• RoleFlow.Flow.Print.emitModule c = (RoleFlow.Flow.Print.emitModuleChecks c).1

def RoleFlow.Flow.Print.permCheckScript (module : String) (checks : Array (String × ℕ × List ℕ)) : String

A script that checks each permutation's numpy expression, applied to 0, …, n - 1, against its index list (as Lean computes it).

Equations

• One or more equations did not get rendered due to their size.

def RoleFlow.Flow.Print.inferBoth {Sg : Sig FS} (what : String) (P : Prog Sg) : Except String ((ρ : List Role) × (ρi : List Role) ×' P.Typed ρ ∧ (P.inverse ρ).Typed ρi)

The roles of P and of its inverse program P.inverse ρ, by inference, with the typings; what names P in the error.

Equations

• One or more equations did not get rendered due to their size.

def RoleFlow.Flow.Print.badPerm? {as bs : List FS} : FOp as bs → Option String

The operation, if it is a permutation that is not PermS.valid (its index map is then not a permutation, and the model does not read it as the steps).

Equations

• One or more equations did not get rendered due to their size.
• RoleFlow.Flow.Print.badPerm? x✝ = none

def RoleFlow.Flow.Print.certify (stages : List Stage) (calls : Bool := false) : Except String Cert

Certify a model: add each stage, after the front end (swapArgs), as a definition of a tower (Linked.add, with its roles and its inverse program's, inferred), then infer the roles of the model as one program and of its inverse program: the stages inlined (modelProg, after the front end), or, with calls, a call of each (callProg). Fails if something is not typed, if two stages have one name, or if a permutation is not valid.

Equations

• One or more equations did not get rendered due to their size.

def RoleFlow.Flow.Print.writeModule (path : System.FilePath) (stages : List Stage) (calls : Bool := false) : IO Unit

Certify a model (certify) and write its module, and <module>_perms.py (permCheckScript).

Equations

• One or more equations did not get rendered due to their size.