The flows' operations over ℝ

The real functions behind the flow signature (RoleFlow.Flow.Basic), on flat vectors Vec n, and the facts its leaves need: the reference measures of the gauss and unit sorts (Lebesgue measure with the standard-normal and the logistic density), change of variables for invertible matrices (through the Jacobian formula), for the elementwise odd powers (off the null set of the axes) and for the changes of measure, and measurability of every conditioner. Ported from SemiInv2; only the weights are now expw of a log-weight.

@[reducible, inline]

abbrev RoleFlow.Flow.Vec (n : ℕ) : Type

A flat vector.

Equations

• RoleFlow.Flow.Vec n = (Fin n → ℝ)

Splitting and joining

def RoleFlow.Flow.vsplit (m n : ℕ) (v : Vec (m + n)) : Vec m × Vec n

Vec (m + n) cut into its first m and last n entries.

Equations

• RoleFlow.Flow.vsplit m n v = (fun (i : Fin m) => v (Fin.castAdd n i), fun (j : Fin n) => v (Fin.natAdd m j))

def RoleFlow.Flow.vjoin {m n : ℕ} (a : Vec m) (b : Vec n) : Vec (m + n)

Concatenation.

Equations

• RoleFlow.Flow.vjoin a b = Fin.append a b

theorem RoleFlow.Flow.vjoin_vsplit (m n : ℕ) (v : Vec (m + n)) : vjoin (vsplit m n v).1 (vsplit m n v).2 = v

@[simp]

theorem RoleFlow.Flow.vsplit_vjoin {m n : ℕ} (a : Vec m) (b : Vec n) : vsplit m n (vjoin a b) = (a, b)

theorem RoleFlow.Flow.measurePreserving_vsplit (m n : ℕ) : MeasureTheory.MeasurePreserving (vsplit m n) MeasureTheory.volume (MeasureTheory.volume.prod MeasureTheory.volume)

theorem RoleFlow.Flow.measurable_vjoin {m n : ℕ} : Measurable fun (p : Vec m × Vec n) => vjoin p.1 p.2

theorem RoleFlow.Flow.measurePreserving_vjoin (m n : ℕ) : MeasureTheory.MeasurePreserving (fun (p : Vec m × Vec n) => vjoin p.1 p.2) (MeasureTheory.volume.prod MeasureTheory.volume) MeasureTheory.volume

Conditioners

def RoleFlow.Flow.natGet {k : ℕ} (v : Vec k) (i : ℕ) : ℝ

A flat read at a ℕ index; 0 out of range.

Equations

• RoleFlow.Flow.natGet v i = if h : i < k then v ⟨i, h⟩ else 0

noncomputable def RoleFlow.Flow.fmvmul (m n : ℕ) (W : Vec (m * n)) (x : Vec n) : Vec m

mvmul: a row-major m × n matrix against a vector.

Equations

• RoleFlow.Flow.fmvmul m n W x i = ∑ j : Fin n, RoleFlow.Flow.natGet W (↑i * n + ↑j) * x j

def RoleFlow.Flow.faddc (c h w : ℕ) (v : Vec (c * (h * w))) (b : Vec c) : Vec (c * (h * w))

addc: a per-channel bias broadcast over h * w pixels.

Equations

• RoleFlow.Flow.faddc c h w v b i = v i + RoleFlow.Flow.natGet b (↑i / (h * w))

noncomputable def RoleFlow.Flow.fvrelu {n : ℕ} (v : Vec n) : Vec n

relu.

Equations

• RoleFlow.Flow.fvrelu v i = max (v i) 0

noncomputable def RoleFlow.Flow.fvtanh {n : ℕ} (v : Vec n) : Vec n

tanh, the bounded log-scale head.

Equations

• RoleFlow.Flow.fvtanh v i = Real.tanh (v i)

noncomputable def RoleFlow.Flow.fconv1x1 (cin cout h w : ℕ) (K : Vec (cout * cin)) (v : Vec (cin * (h * w))) : Vec (cout * (h * w))

conv1x1: per-pixel channel mixing.

Equations

• RoleFlow.Flow.fconv1x1 cin cout h w K v i = ∑ ci : Fin cin, RoleFlow.Flow.natGet K (↑i / (h * w) * cin + ↑ci) * RoleFlow.Flow.natGet v (↑ci * (h * w) + ↑i % (h * w))

noncomputable def RoleFlow.Flow.fconv3x3 (cin cout h w : ℕ) (K : Vec (cout * (cin * 9))) (v : Vec (cin * (h * w))) : Vec (cout * (h * w))

conv3x3: SAME-zero-padded cross-correlation; out-of-image offsets contribute 0.

Equations

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

def RoleFlow.Flow.frows2 (m n : ℕ) (v : Vec ((m + m) * n)) : Vec (m * n) × Vec (m * n)

rows2: the two-row-block split of a matrix.

Equations

• RoleFlow.Flow.frows2 m n v = (fun (i : Fin (m * n)) => RoleFlow.Flow.natGet v ↑i, fun (i : Fin (m * n)) => RoleFlow.Flow.natGet v (m * n + ↑i))

noncomputable def RoleFlow.Flow.flayernorm {n : ℕ} (v : Vec n) : Vec n

Layer norm without gain and bias, with ε = 10⁻⁵.

Equations

• RoleFlow.Flow.flayernorm v = fun (i : Fin n) => (v i - (∑ j : Fin n, v j) / ↑n) / √((∑ j : Fin n, (v j - (∑ j : Fin n, v j) / ↑n) ^ 2) / ↑n + 1 / 100000)

noncomputable def RoleFlow.Flow.fgelu {n : ℕ} (v : Vec n) : Vec n

gelu, in its tanh approximation.

Equations

• RoleFlow.Flow.fgelu v i = v i / 2 * (1 + Real.tanh (√(2 / Real.pi) * (v i + 44715e-6 * v i ^ 3)))

noncomputable def RoleFlow.Flow.fattend (m h : ℕ) (q : Vec h) (K Vs : Vec (h * m)) : Vec h

Single-head dot-product attention: key and value j are the row j * h + ·.

Equations

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

Densities of the reference measures

noncomputable def RoleFlow.Flow.stdNormalPDF (n : ℕ) (x : Vec n) : ℝ

The standard-normal density on Vec n.

Equations

• RoleFlow.Flow.stdNormalPDF n x = ∏ i : Fin n, ProbabilityTheory.gaussianPDFReal 0 1 (x i)

theorem RoleFlow.Flow.stdNormalPDF_pos (n : ℕ) (x : Vec n) : 0 < stdNormalPDF n x

noncomputable def RoleFlow.Flow.logisticPDF (n : ℕ) (t : Vec n) : ℝ

The logistic density σ'(t) = e⁻ᵗ / (1 + e⁻ᵗ)² on Vec n: the density of a point of (0, 1)ⁿ (uniform) in its logit chart.

Equations

• RoleFlow.Flow.logisticPDF n t = ∏ i : Fin n, Real.exp (-t i) / (1 + Real.exp (-t i)) ^ 2

theorem RoleFlow.Flow.logisticPDF_pos (n : ℕ) (t : Vec n) : 0 < logisticPDF n t

theorem RoleFlow.Flow.measurable_stdNormalPDF (n : ℕ) : Measurable (stdNormalPDF n)

theorem RoleFlow.Flow.measurable_logisticPDF (n : ℕ) : Measurable (logisticPDF n)

noncomputable def RoleFlow.Flow.stdNormal (n : ℕ) : MeasureTheory.Measure (Vec n)

The standard normal on Vec n.

Equations

• RoleFlow.Flow.stdNormal n = MeasureTheory.volume.withDensity fun (x : RoleFlow.Flow.Vec n) => ENNReal.ofReal (RoleFlow.Flow.stdNormalPDF n x)

noncomputable def RoleFlow.Flow.stdLogistic (n : ℕ) : MeasureTheory.Measure (Vec n)

The logistic distribution on Vec n.

Equations

• RoleFlow.Flow.stdLogistic n = MeasureTheory.volume.withDensity fun (x : RoleFlow.Flow.Vec n) => ENNReal.ofReal (RoleFlow.Flow.logisticPDF n x)

instance RoleFlow.Flow.instSigmaFiniteVecStdNormal (n : ℕ) : MeasureTheory.SigmaFinite (stdNormal n)

instance RoleFlow.Flow.instSigmaFiniteVecStdLogistic (n : ℕ) : MeasureTheory.SigmaFinite (stdLogistic n)

theorem RoleFlow.Flow.cov_toDensity {n : ℕ} (ρ : Vec n → ℝ) (hm : Measurable ρ) (hp : ∀ (x : Vec n), 0 < ρ x) : Cov MeasureTheory.volume (MeasureTheory.volume.withDensity fun (x : Vec n) => ENNReal.ofReal (ρ x)) (fun (x : Vec n) => x) fun (x : Vec n) => expw (Real.log (ρ x))

From Lebesgue measure to volume.withDensity ρ by the identity, with log-weight log ρ.

Change of variables through the Jacobian

theorem RoleFlow.Flow.cov_of_hasFDerivAt {E : Type u_1} [NormedAddCommGroup E] [NormedSpace ℝ E] [FiniteDimensional ℝ E] [MeasurableSpace E] [BorelSpace E] {μ : MeasureTheory.Measure E} [μ.IsAddHaarMeasure] {f : E → E} {f' : E → E →L[ℝ] E} (hf' : ∀ (x : E), HasFDerivAt f (f' x) x) (hf : Function.Bijective f) : Cov μ μ f fun (x : E) => ENNReal.ofReal |(f' x).det|

A differentiable bijection of a finite-dimensional space with Haar measure satisfies change of variables with |det f'|.

theorem RoleFlow.Flow.cov_mulVec {n : ℕ} (M : Matrix (Fin n) (Fin n) ℝ) (h : IsUnit M.det) : Cov MeasureTheory.volume MeasureTheory.volume M.mulVec fun (x : Fin n → ℝ) => ENNReal.ofReal |M.det|

An invertible matrix satisfies change of variables with |det M|.

theorem RoleFlow.Flow.cov_scale {n : ℕ} (s : Vec n) : Cov MeasureTheory.volume MeasureTheory.volume (fun (x : Vec n) (i : Fin n) => x i * Real.exp (s i)) fun (x : Vec n) => expw (∑ i : Fin n, s i)

x ↦ x ⊙ exp s, with log-weight Σ s.

theorem RoleFlow.Flow.cov_add {n : ℕ} (t : Vec n) : Cov MeasureTheory.volume MeasureTheory.volume (fun (x : Vec n) (i : Fin n) => x i + t i) fun (x : Vec n) => 1

A translation preserves Lebesgue measure.

theorem RoleFlow.Flow.cov_perm {n : ℕ} (σ : Equiv.Perm (Fin n)) : Cov MeasureTheory.volume MeasureTheory.volume (fun (v : Vec n) (i : Fin n) => v (σ i)) fun (x : Vec n) => 1

v ↦ v ∘ σ preserves Lebesgue measure.

The triangular channel map (Glow's LU 1 × 1 convolution)

noncomputable def RoleFlow.Flow.triuMat (c n : ℕ) (U : Vec (c * c)) : Matrix (Fin (c * n)) (Fin (c * n)) ℝ

The matrix of triu c n on Vec (c * n), c channels of n pixels. Entry i is channel i / n at pixel i % n. Within a pixel it is upper triangular, with diagonal exp U_aa and entries U_ab above it (U row-major c × c).

Equations

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

noncomputable def RoleFlow.Flow.ftriu (c n : ℕ) (x : Vec (c * n)) (U : Vec (c * c)) : Vec (c * n)

x ↦ M x.

Equations

• RoleFlow.Flow.ftriu c n x U = (RoleFlow.Flow.triuMat c n U).mulVec x

noncomputable def RoleFlow.Flow.ftriuInv (c n : ℕ) (x : Vec (c * n)) (U : Vec (c * c)) : Vec (c * n)

x ↦ M⁻¹ x, a triangular solve.

Equations

• RoleFlow.Flow.ftriuInv c n x U = (RoleFlow.Flow.triuMat c n U)⁻¹.mulVec x

noncomputable def RoleFlow.Flow.triuLogdet (c n : ℕ) (U : Vec (c * c)) : ℝ

log det M = Σᵢ log Mᵢᵢ = n · Σ_a U_aa.

Equations

• RoleFlow.Flow.triuLogdet c n U = ∑ i : Fin (c * n), RoleFlow.Flow.natGet U (↑i / n * c + ↑i / n)

theorem RoleFlow.Flow.triuMat_blockTriangular {c n : ℕ} (U : Vec (c * c)) : (triuMat c n U).BlockTriangular id

theorem RoleFlow.Flow.triuMat_det {c n : ℕ} (U : Vec (c * c)) : (triuMat c n U).det = Real.exp (triuLogdet c n U)

theorem RoleFlow.Flow.triuMat_isUnit {c n : ℕ} (U : Vec (c * c)) : IsUnit (triuMat c n U).det

theorem RoleFlow.Flow.ftriuInv_ftriu {c n : ℕ} (x : Vec (c * n)) (U : Vec (c * c)) : ftriuInv c n (ftriu c n x U) U = x

theorem RoleFlow.Flow.ftriu_ftriuInv {c n : ℕ} (x : Vec (c * n)) (U : Vec (c * c)) : ftriu c n (ftriuInv c n x U) U = x

theorem RoleFlow.Flow.cov_triu {c n : ℕ} (U : Vec (c * c)) : Cov MeasureTheory.volume MeasureTheory.volume (fun (x : Vec (c * n)) => ftriu c n x U) fun (x : Vec (c * n)) => expw (triuLogdet c n U)

Elementwise odd powers and their roots

x ↦ x ^ (2k + 1) for k : ℕ (the cube at k = 1), a bijection of ℝ whose derivative (2k + 1) x ^ 2k vanishes only at 0.

@[reducible, inline]

abbrev RoleFlow.Flow.oddExp (k : ℕ) : ℕ

The odd exponent 2k + 1.

Equations

• RoleFlow.Flow.oddExp k = 2 * k + 1

theorem RoleFlow.Flow.odd_oddExp (k : ℕ) : Odd (oddExp k)

noncomputable def RoleFlow.Flow.oroot (k : ℕ) (x : ℝ) : ℝ

The real (2k + 1)-th root: x ^ (1/(2k+1)) on x ≥ 0, extended oddly.

Equations

• RoleFlow.Flow.oroot k x = if 0 ≤ x then x ^ (↑(RoleFlow.Flow.oddExp k))⁻¹ else -(-x) ^ (↑(RoleFlow.Flow.oddExp k))⁻¹

theorem RoleFlow.Flow.oroot_pow (k : ℕ) (x : ℝ) : oroot k (x ^ oddExp k) = x

theorem RoleFlow.Flow.pow_oroot (k : ℕ) (y : ℝ) : oroot k y ^ oddExp k = y

def RoleFlow.Flow.fopow (k : ℕ) {n : ℕ} (v : Vec n) : Vec n

x ↦ x ^ (2k + 1), elementwise.

Equations

• RoleFlow.Flow.fopow k v i = v i ^ RoleFlow.Flow.oddExp k

noncomputable def RoleFlow.Flow.foroot (k : ℕ) {n : ℕ} (v : Vec n) : Vec n

y ↦ y ^ (1/(2k + 1)), elementwise.

Equations

• RoleFlow.Flow.foroot k v i = RoleFlow.Flow.oroot k (v i)

theorem RoleFlow.Flow.foroot_fopow (k : ℕ) {n : ℕ} (v : Vec n) : foroot k (fopow k v) = v

theorem RoleFlow.Flow.fopow_foroot (k : ℕ) {n : ℕ} (v : Vec n) : fopow k (foroot k v) = v

noncomputable def RoleFlow.Flow.opowDet1 (k : ℕ) (x : ℝ) : ℝ

(2k + 1) x ^ 2k, or 1 at x = 0: the derivative off the axes, kept positive so that its log is finite everywhere. The axes are null, so change of variables holds with it (cov_fopow).

Equations

• RoleFlow.Flow.opowDet1 k x = if x = 0 then 1 else ↑(RoleFlow.Flow.oddExp k) * x ^ (2 * k)

theorem RoleFlow.Flow.opowDet1_pos (k : ℕ) (x : ℝ) : 0 < opowDet1 k x

noncomputable def RoleFlow.Flow.opowLogdet (k : ℕ) {n : ℕ} (v : Vec n) : ℝ

Σᵢ log ((2k + 1) xᵢ ^ 2k) off the axes: the log-determinant of the elementwise power.

Equations

• RoleFlow.Flow.opowLogdet k v = ∑ i : Fin n, Real.log (RoleFlow.Flow.opowDet1 k (v i))

theorem RoleFlow.Flow.measurable_oroot (k : ℕ) : Measurable (oroot k)

theorem RoleFlow.Flow.measurable_fopow (k n : ℕ) : Measurable (fopow k)

theorem RoleFlow.Flow.measurable_foroot (k n : ℕ) : Measurable (foroot k)

theorem RoleFlow.Flow.measurable_opowDet1 (k : ℕ) : Measurable (opowDet1 k)

theorem RoleFlow.Flow.measurable_opowLogdet (k n : ℕ) : Measurable (opowLogdet k)

theorem RoleFlow.Flow.hasFDerivAt_fopow (k : ℕ) {n : ℕ} (x : Vec n) : HasFDerivAt (fopow k) (LinearMap.toContinuousLinearMap (Matrix.toLin' (Matrix.diagonal fun (i : Fin n) => ↑(oddExp k) * x i ^ (2 * k)))) x

theorem RoleFlow.Flow.ae_ne_zero (n : ℕ) : ∀ᵐ (x : Vec n), ∀ (i : Fin n), x i ≠ 0

The axes are null.

theorem RoleFlow.Flow.cov_fopow (k n : ℕ) : Cov MeasureTheory.volume MeasureTheory.volume (fopow k) fun (x : Vec n) => expw (opowLogdet k x)

x ↦ x ^ (2k + 1), with log-weight opowLogdet k.

Measurability of the conditioners

theorem RoleFlow.Flow.Measurable.natGet {α : Type u_1} [MeasurableSpace α] {k : ℕ} {f : α → Vec k} (hf : Measurable f) (i : ℕ) : Measurable fun (a : α) => Flow.natGet (f a) i

theorem RoleFlow.Flow.measurable_tanh : Measurable Real.tanh

theorem RoleFlow.Flow.Continuous.natGet {α : Type u_1} [TopologicalSpace α] {k : ℕ} {f : α → Vec k} (hf : Continuous f) (i : ℕ) : Continuous fun (a : α) => Flow.natGet (f a) i

theorem RoleFlow.Flow.continuous_triuMat (c n : ℕ) : Continuous fun (U : Vec (c * c)) => triuMat c n U

theorem RoleFlow.Flow.continuous_triuMat_inv (c n : ℕ) : Continuous fun (U : Vec (c * c)) => (triuMat c n U)⁻¹

theorem RoleFlow.Flow.measurable_ftriu (c n : ℕ) : Measurable fun (p : Vec (c * n) × Vec (c * c)) => ftriu c n p.1 p.2

theorem RoleFlow.Flow.measurable_ftriuInv (c n : ℕ) : Measurable fun (p : Vec (c * n) × Vec (c * c)) => ftriuInv c n p.1 p.2

theorem RoleFlow.Flow.measurable_triuLogdet (c n : ℕ) : Measurable (triuLogdet c n)