From afe30280459bd2789f97f3a1517127258d23643d Mon Sep 17 00:00:00 2001 From: Samuel Schlesinger Date: Sat, 3 Oct 2026 23:29:07 -0400 Subject: [PATCH 1/3] feat(Probability): add PMF lemmas for maps, products, and failure bounds Adds further general lemmas to `Cslib.Probability.PMF`: point masses under injective maps and equivalences, lower bounds by a single preimage or execution path, congruence of `bind` and `map` on the support, dependent pairing and the marginals of a pairing bind, point masses of images of a uniform input, and `uniformOfFintype_prod`, which identifies a uniform pair with two independent uniform samples. `toOuterMeasure_bind_failure_le` bounds the failure probability of a randomized continuation: if every possible input satisfying a precondition leads to failure with probability at most `error`, the composite fails with probability at most `P(precondition fails) + error`. It needs no independence or finiteness assumption. `toOuterMeasure_rate_ge` derives the averaging consequence: if a test beats a threshold on average, the excess mass consists of inputs whose conditional acceptance probability reaches the threshold. --- Cslib/Probability/PMF.lean | 129 +++++++++++++++++++++++++++++++++++++ 1 file changed, 129 insertions(+) diff --git a/Cslib/Probability/PMF.lean b/Cslib/Probability/PMF.lean index cdbd7fa7e..28ecd5f29 100644 --- a/Cslib/Probability/PMF.lean +++ b/Cslib/Probability/PMF.lean @@ -28,6 +28,9 @@ the Mathlib module instead. - `Cslib.Probability.PMF.bind_pair_tsum_fst`: marginalizing over the first component - `Cslib.Probability.PMF.uniformOfFintype_map_equiv`: a uniform distribution is invariant under equivalence +- `Cslib.Probability.PMF.uniformOfFintype_prod`: a uniform pair is two independent uniform samples +- `Cslib.Probability.PMF.toOuterMeasure_bind_failure_le`: errors compose through a + state-dependent continuation without independence assumptions - `Cslib.Probability.PMF.posteriorDist`: the posterior as a `PMF` - `Cslib.Probability.PMF.posteriorDist_eq_prior_of_outputIndist`: if the output distribution does not depend on the input, conditioning does @@ -51,6 +54,32 @@ theorem summable_toReal (p : PMF α) : Summable (fun a => (p a).toReal) := @[simp] theorem tsum_toReal (p : PMF α) : ∑' a, (p a).toReal = 1 := by rw [← ENNReal.tsum_toReal_eq p.apply_ne_top, p.tsum_coe, ENNReal.toReal_one] +/-- An injective representation of outcomes also embeds their complete probability laws. -/ +theorem map_injective {f : α → β} (hf : Function.Injective f) : + Function.Injective (PMF.map f) := by + classical + intro p q h + ext a + simpa only [PMF.map_apply, hf.eq_iff, tsum_ite_eq'] using congrArg (fun law => law (f a)) h + +/-- Relabeling a distribution by an equivalence preserves each corresponding point mass. -/ +theorem map_equiv_apply (p : PMF α) (e : α ≃ β) (b : β) : + p.map e b = p (e.symm b) := by + classical + simp only [PMF.map_apply, ← e.symm_apply_eq, tsum_ite_eq'] + +/-- A deterministic map preserves at least the mass of each individual preimage. -/ +theorem le_map_apply (p : PMF α) (f : α → β) (a : α) : p a ≤ (p.map f) (f a) := by + classical + rw [PMF.map_apply] + simpa using ENNReal.le_tsum (f := fun x => if f a = f x then p x else 0) a + +/-- A randomized continuation preserves at least the joint mass of any one execution path. -/ +theorem mul_le_bind_apply (p : PMF α) (f : α → PMF β) (a : α) (b : β) : + p a * f a b ≤ p.bind f b := by + rw [PMF.bind_apply] + exact ENNReal.le_tsum (f := fun x => p x * f x b) a + /-- The real-valued probabilities of a finite distribution sum to one. -/ theorem sum_toReal [Fintype α] (p : PMF α) : ∑ a, (p a).toReal = 1 := by simpa using tsum_toReal p @@ -108,6 +137,47 @@ theorem toOuterMeasure_bind_toReal [Fintype α] (p : PMF α) (kernel : α → PM (fun a _ => ENNReal.mul_ne_top (PMF.apply_ne_top _ _) (toOuterMeasure_ne_top _ _))] simp +/-- A randomized continuation can fail because its precondition was already false or because +it fails from a valid input. No independence or finite-state assumption is needed. -/ +theorem toOuterMeasure_bind_failure_le (p : PMF α) (kernel : α → PMF β) + (pre : Set α) (failure : Set β) (error : ENNReal) + (hstep : ∀ a ∈ p.support, a ∈ pre → (kernel a).toOuterMeasure failure ≤ error) : + (p.bind kernel).toOuterMeasure failure ≤ p.toOuterMeasure preᶜ + error := by + rw [PMF.toOuterMeasure_bind_apply, PMF.toOuterMeasure_apply, ← one_mul error, ← p.tsum_coe, + ← ENNReal.tsum_mul_right, ← ENNReal.tsum_add] + refine ENNReal.tsum_le_tsum fun a => ?_ + by_cases ha : a ∈ pre + · rw [Set.indicator_of_notMem (by simpa using ha), zero_add] + by_cases hs : a ∈ p.support + · exact mul_le_mul_of_nonneg_left (hstep a hs ha) zero_le + · simp [(PMF.apply_eq_zero_iff p a).2 hs] + · rw [Set.indicator_of_mem ha] + exact le_add_right (mul_le_of_le_one_right' (toOuterMeasure_le_one _ _)) + +/-- The real-valued error rule for a randomized continuation. -/ +theorem toOuterMeasure_bind_failure_toReal_le (p : PMF α) (kernel : α → PMF β) + (pre : Set α) (failure : Set β) {error : ℝ} (herror : 0 ≤ error) + (hstep : ∀ a ∈ p.support, a ∈ pre → + ((kernel a).toOuterMeasure failure).toReal ≤ error) : + ((p.bind kernel).toOuterMeasure failure).toReal ≤ + (p.toOuterMeasure preᶜ).toReal + error := by + rw [← ENNReal.toReal_ofReal herror, ← ENNReal.toReal_add (toOuterMeasure_ne_top _ _) + ofReal_ne_top] + exact ENNReal.toReal_mono (add_ne_top.2 ⟨toOuterMeasure_ne_top _ _, ofReal_ne_top⟩) + (toOuterMeasure_bind_failure_le p kernel pre failure _ fun a ha hpre => + (le_ofReal_iff_toReal_le (toOuterMeasure_ne_top _ _) herror).2 (hstep a ha hpre)) + +/-- If a randomized test beats a threshold on average, at least its excess probability mass +consists of inputs whose conditional acceptance probability reaches that threshold. -/ +theorem toOuterMeasure_rate_ge (p : PMF α) (kernel : α → PMF β) (event : Set β) + {threshold : ℝ} (hthreshold : 0 ≤ threshold) : + ((p.bind kernel).toOuterMeasure event).toReal - threshold ≤ + (p.toOuterMeasure {a | threshold ≤ ((kernel a).toOuterMeasure event).toReal}).toReal := by + have h := toOuterMeasure_bind_failure_toReal_le p kernel + {a | ((kernel a).toOuterMeasure event).toReal < threshold} event hthreshold + (fun a _ ha => le_of_lt ha) + simpa only [Set.compl_ofPred, not_lt] using (sub_le_iff_le_add.mpr h) + /-- Randomized postprocessing averages the outcome probabilities over any discrete input. -/ theorem bind_apply_toReal_tsum (p : PMF α) (kernel : α → PMF β) (b : β) : (p.bind kernel b).toReal = ∑' a, (p a).toReal * (kernel a b).toReal := by @@ -162,6 +232,29 @@ theorem sum_mul_le [Fintype α] (p : PMF α) (score : α → ℝ) (bound : ℝ) mul_le_mul_of_nonneg_left (hscore a) ENNReal.toReal_nonneg) _ = bound := by rw [← Finset.sum_mul, sum_toReal, one_mul] +/-- The mass of an image of a uniform input is its fiber size divided by the input size. -/ +theorem uniformOfFintype_map_apply [Fintype α] [Nonempty α] (f : α → β) (b : β) : + ((PMF.uniformOfFintype α).map f) b = + (Nat.card {a // f a = b} : ℝ≥0∞) / Fintype.card α := by + classical + simp only [PMF.map_apply, PMF.uniformOfFintype_apply, tsum_fintype, + ← Finset.sum_filter, Finset.sum_const, nsmul_eq_mul, div_eq_mul_inv, + Nat.card_eq_fintype_card, Fintype.card_subtype, eq_comm] + +/-- A continuation only needs to agree on outcomes which the preceding distribution can produce. -/ +theorem bind_congr_on_support (p : PMF α) (f g : α → PMF β) + (h : ∀ a ∈ p.support, f a = g a) : p.bind f = p.bind g := by + ext b + refine tsum_congr fun a => ?_ + by_cases ha : a ∈ p.support + · rw [h a ha] + · simp [(PMF.apply_eq_zero_iff p a).2 ha] + +/-- Re-encoding only needs to agree on outcomes which the distribution can produce. -/ +theorem map_congr_on_support (p : PMF α) (f g : α → β) + (h : ∀ a ∈ p.support, f a = g a) : p.map f = p.map g := + bind_congr_on_support p _ _ (fun a ha => congrArg PMF.pure (h a ha)) + /-- Evaluating the "pairing" bind `(do let a ← p; return (a, ← f a))` at `(a, b)` gives the product `p a * f a b`. -/ theorem bind_pair_apply (p : PMF α) (f : α → PMF β) (a : α) (b : β) : @@ -172,12 +265,39 @@ theorem bind_pair_apply (p : PMF α) (f : α → PMF β) (a : α) (b : β) : · intro b' hb'; simp [PMF.pure_apply, hb'.symm] · intro a' ha'; rw [PMF.bind_apply]; simp [PMF.pure_apply, ha'.symm] +/-- The pairing law also holds when the second sample's type depends on the first. -/ +theorem bind_sigma_apply {β : α → Type*} (p : PMF α) (f : (a : α) → PMF (β a)) + (a : α) (b : β a) : + (p.bind (fun a => (f a).map (Sigma.mk a))) ⟨a, b⟩ = p a * f a b := by + classical + rw [PMF.bind_apply, tsum_eq_single a] + · congr 1 + simp [PMF.map_apply] + · intro i hi + have hne (c : β i) : (Sigma.mk a b : Sigma β) ≠ ⟨i, c⟩ := + fun h => hi (congrArg Sigma.fst h).symm + simp [PMF.map_apply, hne] + /-- Summing the pairing bind over the first component gives the marginal. -/ theorem bind_pair_tsum_fst (p : PMF α) (f : α → PMF β) (b : β) : ∑' a, (p.bind fun a' => (f a').bind fun b' => PMF.pure (a', b')) (a, b) = (p.bind f) b := by simp_rw [bind_pair_apply, PMF.bind_apply] +/-- Marginalizing a joint distribution over its freshly sampled second component. -/ +@[simp] theorem map_fst_bind_pair (p : PMF α) (f : α → PMF β) : + (p.bind (fun a => (f a).map (a, ·))).map Prod.fst = p := by + simp only [PMF.map_bind, PMF.map_comp, Function.comp_def] + change p.bind (fun a => (f a).map (Function.const β a)) = p + simp only [PMF.map_const, PMF.bind_pure] + +/-- Marginalizing a joint distribution over its first component gives ordinary sequencing. -/ +@[simp] theorem map_snd_bind_pair (p : PMF α) (f : α → PMF β) : + (p.bind (fun a => (f a).map (a, ·))).map Prod.snd = p.bind f := by + simp only [PMF.map_bind, PMF.map_comp, Function.comp_def] + change p.bind (fun a => (f a).map id) = p.bind f + simp only [PMF.map_id] + /-- A uniform distribution on a finite type is invariant under any equivalence. -/ theorem uniformOfFintype_map_equiv {γ : Type v} [Fintype α] [Fintype γ] [Nonempty α] [Nonempty γ] (e : α ≃ γ) : @@ -187,6 +307,15 @@ theorem uniformOfFintype_map_equiv {γ : Type v} [Fintype α] [Fintype γ] [None · simp [Fintype.card_congr e] · exact fun a ha => ite_eq_right fun h => ha (by simp [h]) +/-- A uniform pair consists of two independent uniform samples. -/ +theorem uniformOfFintype_prod [Fintype α] [Fintype β] [Nonempty α] [Nonempty β] : + PMF.uniformOfFintype (α × β) = (PMF.uniformOfFintype α).bind + (fun a => (PMF.uniformOfFintype β).map (fun b => (a, b))) := by + ext ⟨a, b⟩ + simp only [PMF.map, Function.comp_def, bind_pair_apply, PMF.uniformOfFintype_apply, + Fintype.card_prod, Nat.cast_mul] + rw [ENNReal.mul_inv] <;> simp + /-- The posterior distribution `Pr[A = a | B = b]` as a `PMF`, given `a ← p`, `b ← f a`, and that `b` has positive marginal probability: the joint distribution's slice at `b`, normalized. -/ From 2e0b6d900f1c85828835161df7086ccc7bcfe727 Mon Sep 17 00:00:00 2001 From: Samuel Schlesinger Date: Sun, 4 Oct 2026 19:01:53 -0400 Subject: [PATCH 2/3] style(Probability): unsqueeze terminal simps in PMF map and product lemmas --- Cslib/Probability/PMF.lean | 18 ++++++++---------- 1 file changed, 8 insertions(+), 10 deletions(-) diff --git a/Cslib/Probability/PMF.lean b/Cslib/Probability/PMF.lean index 28ecd5f29..7216ae272 100644 --- a/Cslib/Probability/PMF.lean +++ b/Cslib/Probability/PMF.lean @@ -60,13 +60,13 @@ theorem map_injective {f : α → β} (hf : Function.Injective f) : classical intro p q h ext a - simpa only [PMF.map_apply, hf.eq_iff, tsum_ite_eq'] using congrArg (fun law => law (f a)) h + simpa [hf.eq_iff] using congrArg (fun law => law (f a)) h /-- Relabeling a distribution by an equivalence preserves each corresponding point mass. -/ theorem map_equiv_apply (p : PMF α) (e : α ≃ β) (b : β) : p.map e b = p (e.symm b) := by classical - simp only [PMF.map_apply, ← e.symm_apply_eq, tsum_ite_eq'] + simp [← e.symm_apply_eq] /-- A deterministic map preserves at least the mass of each individual preimage. -/ theorem le_map_apply (p : PMF α) (f : α → β) (a : α) : p a ≤ (p.map f) (f a) := by @@ -176,7 +176,7 @@ theorem toOuterMeasure_rate_ge (p : PMF α) (kernel : α → PMF β) (event : Se have h := toOuterMeasure_bind_failure_toReal_le p kernel {a | ((kernel a).toOuterMeasure event).toReal < threshold} event hthreshold (fun a _ ha => le_of_lt ha) - simpa only [Set.compl_ofPred, not_lt] using (sub_le_iff_le_add.mpr h) + simpa [Set.compl_ofPred] using h /-- Randomized postprocessing averages the outcome probabilities over any discrete input. -/ theorem bind_apply_toReal_tsum (p : PMF α) (kernel : α → PMF β) (b : β) : @@ -237,9 +237,7 @@ theorem uniformOfFintype_map_apply [Fintype α] [Nonempty α] (f : α → β) (b ((PMF.uniformOfFintype α).map f) b = (Nat.card {a // f a = b} : ℝ≥0∞) / Fintype.card α := by classical - simp only [PMF.map_apply, PMF.uniformOfFintype_apply, tsum_fintype, - ← Finset.sum_filter, Finset.sum_const, nsmul_eq_mul, div_eq_mul_inv, - Nat.card_eq_fintype_card, Fintype.card_subtype, eq_comm] + simp [← Finset.sum_filter, div_eq_mul_inv, Fintype.card_subtype, eq_comm] /-- A continuation only needs to agree on outcomes which the preceding distribution can produce. -/ theorem bind_congr_on_support (p : PMF α) (f g : α → PMF β) @@ -272,11 +270,11 @@ theorem bind_sigma_apply {β : α → Type*} (p : PMF α) (f : (a : α) → PMF classical rw [PMF.bind_apply, tsum_eq_single a] · congr 1 - simp [PMF.map_apply] + simp · intro i hi have hne (c : β i) : (Sigma.mk a b : Sigma β) ≠ ⟨i, c⟩ := fun h => hi (congrArg Sigma.fst h).symm - simp [PMF.map_apply, hne] + simp [hne] /-- Summing the pairing bind over the first component gives the marginal. -/ theorem bind_pair_tsum_fst (p : PMF α) (f : α → PMF β) (b : β) : @@ -289,14 +287,14 @@ theorem bind_pair_tsum_fst (p : PMF α) (f : α → PMF β) (b : β) : (p.bind (fun a => (f a).map (a, ·))).map Prod.fst = p := by simp only [PMF.map_bind, PMF.map_comp, Function.comp_def] change p.bind (fun a => (f a).map (Function.const β a)) = p - simp only [PMF.map_const, PMF.bind_pure] + simp /-- Marginalizing a joint distribution over its first component gives ordinary sequencing. -/ @[simp] theorem map_snd_bind_pair (p : PMF α) (f : α → PMF β) : (p.bind (fun a => (f a).map (a, ·))).map Prod.snd = p.bind f := by simp only [PMF.map_bind, PMF.map_comp, Function.comp_def] change p.bind (fun a => (f a).map id) = p.bind f - simp only [PMF.map_id] + simp [PMF.map_id] /-- A uniform distribution on a finite type is invariant under any equivalence. -/ theorem uniformOfFintype_map_equiv {γ : Type v} [Fintype α] [Fintype γ] [Nonempty α] [Nonempty γ] From 4b3ff93eafb49df8c380751e5e36e888d7904550 Mon Sep 17 00:00:00 2001 From: Samuel Schlesinger Date: Sun, 4 Oct 2026 19:18:24 -0400 Subject: [PATCH 3/3] style(Probability): use DFunLike.congr_fun in map_injective --- Cslib/Probability/PMF.lean | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Cslib/Probability/PMF.lean b/Cslib/Probability/PMF.lean index 7216ae272..ec639cf9b 100644 --- a/Cslib/Probability/PMF.lean +++ b/Cslib/Probability/PMF.lean @@ -60,7 +60,7 @@ theorem map_injective {f : α → β} (hf : Function.Injective f) : classical intro p q h ext a - simpa [hf.eq_iff] using congrArg (fun law => law (f a)) h + simpa [hf.eq_iff] using DFunLike.congr_fun h (f a) /-- Relabeling a distribution by an equivalence preserves each corresponding point mass. -/ theorem map_equiv_apply (p : PMF α) (e : α ≃ β) (b : β) :