PMF monadic theorems and `do`notation

Stream: mathlib4

[Yannick Seurin]

I've been playing a bit with PMFs and one thing I noticed is that theorems about the monadic structure of PMFs are not usable when using the do notation (or >>=). Here is an example:

import Mathlib

#check PMF.bind_apply
/-
PMF.bind_apply.{u_1, u_2} {α : Type u_1} {β : Type u_2} (p : PMF α) (f : α → PMF β) (b : β) :
  (p.bind f) b = ∑' (a : α), p a * (f a) b
-/

-- check definitional equality
theorem PMF.bind_eq_do_notation {α β : Type u} (p : PMF α) (f : α → PMF β) (b : β) :
     (p.bind f) b = (do
                      let a ← p
                      f a) b := rfl

theorem bind_skip_const {α β : Type u} (pa : PMF α) (pb : PMF β)
    (f : α → PMF β) (h : ∀ a : α, f a = pb) :
    (do
      let a ← pa
      f a) = pb := by
  ext b
  rw [PMF.bind_apply pa f b]
  /-
  Tactic `rewrite` failed: Did not find an occurrence of the pattern
    (pa.bind f) b
  in the target expression
    (do
          let a ← pa
          f a)
        b =
      pb b
  -/
  sorry

My understanding is that p.bind f and the corresponding do notation are not syntactically equal, which in turns is because do and >>= are notations for Bind.bind, not PMF.bind. Is that correct?
The only fix I could come with was to rewrite all PMF monadic theorems with the donotation (see below), but I was wondering whether there is a more clever solution.

import Mathlib

theorem PMF.bind_apply' {α β : Type u} (p : PMF α) (f : α → PMF β) (b : β) :
    (do
      let a ← p
      f a) b = ∑' (a : α), p a * (f a) b :=
  PMF.bind_apply p f b

theorem bind_skip_const' {α β : Type u} (pa : PMF α) (pb : PMF β)
    (f : α → PMF β) (h : ∀ a : α, f a = pb) :
    (do
      let a ← pa
      f a) = pb := by
  ext b
  simp only [PMF.bind_apply', h]
  rw [ENNReal.tsum_mul_right, PMF.tsum_coe pa, one_mul]

[Jovan Gerbscheid]

The problem is that the lemmas are about PMF.bind, but do notation uses Bind.bind.

[Jovan Gerbscheid]

So yeah, I'm afraid you'll have to choose between re-doing all of the lemmas using Bind.bind, or just not using do notation.

[Yannick Seurin]

Thanks! Writing involved random processes without the do notation is pretty cumbersome, so I'll stick with rewritten lemmas!

[Jovan Gerbscheid]

In that case, you might also be able to convince people that the lemmas in mathlib should use Bind.bind instead of PMF.bind. Though it feels kind of funny to me to use do notation for pure maths instead of programming :upside_down:

[Aaron Liu]

is there not a lemme that says PMF.bind is the same as Bind.bind?

[Eric Wieser]

It's not stateable universe polymorphically

[Eric Wieser]

Which is to say that we would need to duplicate all the lemmas, as simply switching them over would lose generality