mk_all

Author: RemyDegenne

.github/workflows/build.yml

@@ -47,7 +47,7 @@ jobs:
           # the single fetched commit.
           fetch-depth: 0 # Fetch all history for all branches and tags
 
-      - name: Build and lint project
+      - name: Build, lint and test project
         uses: leanprover/lean-action@38fbc41a8c28c4cbaec22d7f7de508ec2e7c0dd9 # v1.5.0
         with:
           build: true

RandomDo.lean

@@ -2,7 +2,6 @@ module  -- shake: keep-all --deprecated_module: ignore
 
 public import RandomDo.ForMathlib.MeasureTheory.MeasurableSpace.Embedding
 public import RandomDo.Measurable
-public import RandomDo.Monad.Examples
 public import RandomDo.Monad.ForInInstances
 public import RandomDo.Monad.Instances
 public import RandomDo.Monad.MeasurableSpace
@@ -16,8 +15,8 @@ public import RandomDo.Probability.Tactic
 public import RandomDo.Probability.Thompson
 public import RandomDo.Probability.Trace
 public import RandomDo.Probability.Transfer
+public import RandomDo.Tactic.Deriving
 public import RandomDo.Tactic.Elab
-public import RandomDo.Tactic.Examples
 public import RandomDo.Tactic.ForInStep
 public import RandomDo.Tactic.IsMarkov
 public import RandomDo.Tactic.Lemmas

RandomDo/Measurable.lean

@@ -31,6 +31,7 @@ open MeasureTheory Set
 
 variable {α : Type*} [MeasurableSpace α]
 
+/-- Lists are measurably equivalent to the sigma type of tuples of a given length. -/
 def List.measurableEquivSigmaTuple : List α ≃ᵐ Σ n, Fin n → α where
   toFun := List.equivSigmaTuple
   invFun := List.equivSigmaTuple.symm
@@ -102,6 +103,7 @@ lemma Vector.measurableSpace_eq_comap {n : ℕ} :
           (MeasurableSpace.comap List.ofFn inferInstance) := MeasurableSpace.comap_comp.symm
     _ = _ := by rw [(measurableEmbedding_ofFn n).comap_eq]
 
+/-- Vectors are equivalent to tuples of a given length. -/
 def Vector.measurableEquivTuple {n : ℕ} : Vector α n ≃ᵐ (Fin n → α) where
   toFun v := fun i ↦ v[i]
   invFun := .ofFn
@@ -115,7 +117,8 @@ def Vector.measurableEquivTuple {n : ℕ} : Vector α n ≃ᵐ (Fin n → α) wh
     ext
     simp
 
-instance : MeasurableSpace (Option α) := MeasurableSpace.map some inferInstance
+instance instMeasurableSpaceOption : MeasurableSpace (Option α) :=
+  MeasurableSpace.map some inferInstance
 
 theorem measurableSet_option_iff {s : Set (Option α)} :
     MeasurableSet s ↔ MeasurableSet (some ⁻¹' s) := Iff.rfl

RandomDo/Monad/Examples.lean

@@ -20,19 +20,26 @@ section MeasurableSpace
 
 variable {α : Type*} [MeasurableSpace α]
 
-instance : MeasurableSpace (List α) :=
+instance instMeasurableSpaceList : MeasurableSpace (List α) :=
   MeasurableSpace.comap List.equivSigmaTuple inferInstance
 
-instance : MeasurableSpace (Array α) :=
+instance instMeasurableSpaceArray : MeasurableSpace (Array α) :=
   MeasurableSpace.comap Array.toList inferInstance
 
-instance {n : ℕ} : MeasurableSpace (Vector α n) :=
+instance instMeasurableSpaceVector {n : ℕ} : MeasurableSpace (Vector α n) :=
   MeasurableSpace.comap Vector.toArray inferInstance
 
+instance instMeasurableSpaceSubarray : MeasurableSpace (Subarray α) :=
+  MeasurableSpace.comap (fun s : Subarray α ↦ s.toList) inferInstance
+
 @[fun_prop]
 lemma measurable_toList : Measurable (Array.toList : Array α → List α) :=
   Measurable.of_comap_le fun _ a ↦ a
 
+@[fun_prop]
+lemma measurable_subarray_toList : Measurable (fun s : Subarray α ↦ s.toList) :=
+  Measurable.of_comap_le fun _ a ↦ a
+
 @[fun_prop]
 lemma measurable_toArray {n : ℕ} : Measurable (Vector.toArray : Vector α n → Array α) :=
   Measurable.of_comap_le fun _ a ↦ a
@@ -52,7 +59,7 @@ section Array
   {β : Type u} [mβ : MeasurableSpace β]
   (as : Array α) (b : β) (f : (a : α) → a ∈ as → β → m (ForInStep β)) : m β :=
   let sz := as.usize
-  let rec @[specialize] loop (i : USize) (b : β) : m β := rdo
+  let rec @[specialize, nolint docBlame] loop (i : USize) (b : β) : m β := rdo
     if i < sz then
       let a := as.uget i lcProof
       match (← f a lcProof b) with
@@ -67,7 +74,7 @@ section Array
 protected def Array.measurableSpaceForIn' [MeasurableSpaceMonad m]
   {β : Type u} [mβ : MeasurableSpace β]
   (as : Array α) (b : β) (f : (a : α) → a ∈ as → β → m (ForInStep β)) : m β :=
-  let rec loop (i : Nat) (h : i ≤ as.size) (b : β) : m β := rdo
+  let rec @[nolint docBlame] loop (i : Nat) (h : i ≤ as.size) (b : β) : m β := rdo
     match i, h with
     | 0,   _ => mPure b
     | i+1, h =>
@@ -97,7 +104,7 @@ variable {α β : Type*} [MeasurableSpace α] [MeasurableSpace β] [Ring α]
 protected def List.measurableSpaceForIn' [MeasurableSpaceMonad m]
     {β : Type u} [mβ : MeasurableSpace β] (as : @& List α) (init : β)
     (f : (a : α) → a ∈ as → β → m (ForInStep β)) : m β :=
-  let rec @[specialize]
+  let rec @[specialize, nolint docBlame]
     loop : (as' : @& List α) → (b : β) → Exists (fun bs => bs ++ as' = as) → m β
       | [], b, _    => mPure b
       | a::as', b, h => rdo

RandomDo/Monad/ForInInstances.lean

@@ -21,7 +21,7 @@ universe u v w
 
 /-- A (core) monad automatically defines a (not necessarily lawful) measurable space monad by
 forgetting the measurable space argument. -/
-def Monad.toMeasurableSpaceMonad (m : Type u → Type v) [Monad m] (α : Type u) [MeasurableSpace α] :
+def Monad.toMeasurableSpaceMonad (m : Type u → Type v) (α : Type u) [_mα : MeasurableSpace α] :
     Type v := m α
 
 instance {m : Type u → Type v} [Monad m] :
@@ -54,11 +54,15 @@ section RandomM
 
 open Function
 
+/-- A monad for random number generation. -/
 structure RandomM (Ω : Type w) [MeasurableSpace Ω] (P : Measure Ω)
     (α : Type u) [MeasurableSpace α] where
+  /-- Draws a value from a state of the source of randomness, and hands back the state left for the
+  next draw. -/
   sample : Ω → α × Ω
   measurePreserving : MeasurePreserving sample P ((Measure.map (Prod.fst ∘ sample) P).prod P)
 
+/-- TODO -/
 abbrev SampleM (Ω : Type w) [MeasurableSpace Ω] (P : Measure Ω) :=
   RandomM (ℕ → Ω) (Measure.infinitePi fun _ : ℕ ↦ P)
 

RandomDo/Monad/Instances.lean

@@ -65,6 +65,8 @@ The `mPure` function is overloaded via `MeasurableSpacePure` instances.
 `MeasurableSpacePure` is typically accessed via `MeasurableSpaceMonad` instances, which extend it.
 -/
 class MeasurableSpacePure (f : (α : Type u) → [MeasurableSpace α] → Type v) where
+  /-- Given `a : α` where `α` has a `MeasurableSpace` instance, `mPure a : f α` represents an
+  action that does nothing and returns a -/
   mPure {α : Type u} [MeasurableSpace α] : α → f α
 
 /--
@@ -140,8 +142,6 @@ variable {f : (α : Type u) → [MeasurableSpace α] → Type v} [MeasurableSpac
     g₁ <$>ₘ g₀ <$>ₘ x = (fun a => g₁ (g₀ a)) <$>ₘ x :=
   (comp_mMap x hg₀ hg₁).symm
 
-@[simp] theorem mMap_unit {a : f PUnit} : (fun _ => PUnit.unit) <$>ₘ a = a := by simp
-
 open MeasurableSpaceBind MeasurableSpacePure MeasurableSpaceFunctor
 
 /-- A `MeasurableSpaceMonad` satisfies the measurable space monad laws. -/

RandomDo/Monad/MeasurableSpace.lean

@@ -63,6 +63,7 @@ def randOps : DoOps := { DoOps.default with
     return mkApp2 m α σ
   }
 
+/-- The `do` notation for writing monadic programs depending on a `MeasurableSpace` instance. -/
 syntax (name := randKind) "rdo" doSeq : term
 
 /-- Define `rdo` notation elaborator. -/
@@ -138,7 +139,7 @@ def rdoForDecl := leading_parser
           | none => break
           | some ($y, s') =>
             $s:ident := s'
-            rdo $body)
+            do $body)
     doElems := doElems.push (← `(doSeqItem| for%$tk $[$h? : ]? $x:ident in $xs rdo $body))
     `(doElem| do $doElems*)
   | _ => Macro.throwUnsupported

RandomDo/Monad/Notation.lean

@@ -6,7 +6,8 @@ Authors: Rémy Degenne
 module
 
 public import RandomDo.Probability.Tactic
-public import RandomDo.Tactic.Examples
+public import RandomDo.Tactic.Elab
+public import LeanMachineLearning.ForMathlib.MeasureTheory.Order.MeasurableArg
 
 set_option linter.style.header false
 
@@ -56,6 +57,8 @@ namespace RDo.Thompson
 
 variable {K n : ℕ}
 
+attribute[fun_prop] Measurable.ite
+
 /-! ## The two stages -/
 
 /-- Stage one: fold the history into the per-arm pull counts `N` (started at one) and reward
@@ -96,6 +99,20 @@ instance : IsMarkovKernel (sampleK (K := K)) := by unfold sampleK; infer_instanc
 
 @[simp] lemma sampleK_apply (NS : (Fin K → ℝ) × (Fin K → ℝ)) : sampleK NS = sample NS := rfl
 
+def thompson {K n : ℕ} (hK : 0 < K) (hist : Vector (Fin K × ℝ) n) :
+    Measure (Fin K) := rdo
+  let mut N : Fin K → ℝ := fun _ ↦ 1
+  let mut S : Fin K → ℝ := fun _ ↦ 0
+  for (a, r) in hist rdo
+    N := fun j ↦ if j = a then N j + 1 else N j
+    S := fun j ↦ if j = a then S j + r else S j
+  let mut θ : Fin K → ℝ := fun _ ↦ 0
+  for j in List.finRange K rdo
+    let z ← gaussianReal (S j / N j) (Real.toNNReal (1 / N j))
+    θ := fun k ↦ if k = j then z else θ k
+  have : Nonempty (Fin K) := Fin.pos_iff_nonempty.mp hK
+  return argmax θ
+
 /-- `thompson` is exactly: fold the history into `(N, S)`, draw the posterior sample `θ` given
 them, play `argmax θ`. Both sides elaborate to the same two loops; all that separates them is the
 `return` at the end of each stage. -/

RandomDo/Probability/Thompson.lean

@@ -0,0 +1,79 @@
+/-
+Copyright (c) 2026 Rémy Degenne. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Rémy Degenne
+-/
+module
+
+public import RandomDo.Tactic.Elab
+
+/-!
+# The `@[is_markov]` attribute
+
+Writing an `rdo` program and then stating that it is a Markov kernel are two separate steps, and
+the second is mechanical: it is what `is_markov` does. This attribute runs it at the declaration,
+so a program carries its own instance:
+
+```
+@[is_markov]
+noncomputable def centred (c : ℝ) : Measure ℝ := rdo
+  let x ← gaussianReal c 1
+  return x
+```
+
+## Which statement is generated
+
+A program's *last* argument is read as the kernel's parameter when it is explicit, giving
+`IsMarkov`. Otherwise the program denotes one fixed distribution and the statement is
+`IsProbabilityMeasure`. So `centred` above yields `IsMarkov centred`, while a parameterless program
+yields `IsProbabilityMeasure` of it, and a program whose trailing arguments are instance-implicit —
+`(μ : Measure ℝ) [IsProbabilityMeasure μ]` — yields `IsProbabilityMeasure` of it too, with those
+arguments bound.
+-/
+
+public meta section
+
+open Lean Meta Elab Term MeasureTheory
+
+namespace RDo.Tactic
+
+/-- The statement to prove for `declName`, as described in the module docstring. -/
+def isMarkovStatement (declName : Name) : MetaM Expr := do
+  let info ← getConstInfo declName
+  forallTelescope info.type fun args body ↦ do
+    unless body.isAppOfArity ``MeasureTheory.Measure 2 do
+      throwError "`IsMarkov` can only be derived for a declaration valued in `Measure`, but \
+        {declName} is valued in{indentExpr body}"
+    let f := mkAppN (mkConst declName (info.levelParams.map .param)) args
+    if h : 0 < args.size then
+      let last := args[args.size - 1]
+      if (← last.fvarId!.getDecl).binderInfo.isExplicit then
+        return ← mkForallFVars args.pop (← mkAppM ``IsMarkov #[← mkLambdaFVars #[last] f])
+    mkForallFVars args (← mkAppM ``MeasureTheory.IsProbabilityMeasure #[f])
+
+/-- Prove `isMarkovStatement declName` with the `is_markov` tactic and register it as an
+instance. -/
+def addIsMarkovInstance (declName : Name) : TermElabM Unit := do
+  let goal ← isMarkovStatement declName
+  let proof ← Term.elabTerm (← `(by intros; is_markov)) (some goal)
+  Term.synthesizeSyntheticMVarsNoPostponing
+  let info ← getConstInfo declName
+  let instName := declName ++ `isMarkov
+  addDecl (.thmDecl { name := instName, levelParams := info.levelParams, type := goal,
+                      value := ← instantiateMVars proof })
+  Meta.addInstance instName .global 1000
+
+/-- The `@[is_markov]` attribute. -/
+initialize registerBuiltinAttribute {
+  name := `is_markov
+  descr := "prove that this `rdo` program is a Markov kernel, and register it as an instance"
+  applicationTime := .afterCompilation
+  add := fun declName _stx kind ↦ do
+    unless kind == AttributeKind.global do
+      throwError "`is_markov` must be a global attribute"
+    (addIsMarkovInstance declName).run'.run'
+}
+
+end RDo.Tactic
+
+end