Probability Monads in Cubical Agda

Posted on April 17, 2019
Tags: ,

Cubical Agda has just come out, and I’ve been playing around with it for a bit. There’s a bunch of info out there on the theory of cubical types, and Homotopy Type Theory more generally (cubical type theory is kind of like an “implementation” of Homotopy type theory), but I wanted to make a post demonstrating cubical Agda in practice, and one of its cool uses from a programming perspective.

So What is Cubical Agda?

I don’t really know! Cubical type theory is quite complex (even for a type theory), and I’m not nearly qualified to properly explain it. In lieu of a proper first-principles explanation, then, I’ll try and give a few examples of how it differs from normal Agda, before moving on to the main example of this post.

Imports
{-# OPTIONS --cubical #-}

open import ProbabilityModule.Semirings

module ProbabilityModule.Monad {s} (rng : Semiring s) where

open import Cubical.Core.Everything
open import Cubical.Relation.Everything
open import Cubical.Foundations.Prelude hiding (_≡⟨_⟩_) renaming (_∙_ to _;_)
open import Cubical.HITs.SetTruncation
open import ProbabilityModule.Utils
Extensionality
One of the big annoyances in standard Agda is that we can’t prove the following:
extensionality :  {f g : A  B}
            (∀ x  f x  g x)
            f  g
It’s emblematic of a wider problem in Agda: we can’t say “two things are equal if they always behave the same”. Infinite types, for instance (like streams) are often only equal via bisimulation: we can’t translate this into normal equality in standard Agda. Cubical type theory, though, has a different notion of “equality”, which allow a wide variety of things (including bisimulations and extensional proofs) to be translated into a proper equality
extensionality = funExt
Isomorphisms
One of these such things we can promote to a “proper equality” is an isomorphism. In the cubical repo this is used to prove things about binary numbers: by proving that there’s an isomorphism between the Peano numbers and binary numbers, they can lift any properties on the Peano numbers to the binary numbers.

So those are two useful examples, but the most interesting use I’ve seen so far is the following:

Higher Inductive Types

Higher Inductive Types are an extension of normal inductive types, like the list:
module NormalList where
 data List {a} (A : Set a) : Set a where
   [] : List A
   _∷_ : A  List A  List A

They allow us to add new equations to a type, as well as constructors. To demonstrate what this means, as well as why you’d want it, I’m going to talk about free objects.

Very informally, a free object on some algebra is the minimal type which satisfies the laws of the algebra. Lists, for instance, are the free monoid. They satisfy all of the monoid laws (\bullet is ++ and ϵ\epsilon is []):

(xy)z=x(yz)(x \bullet y) \bullet z = x \bullet (y \bullet z) xϵ=xx \bullet \epsilon = x ϵx=x\epsilon \bullet x = x

But nothing else. That means they don’t satisfy any extra laws (like, for example, commutativity), and they don’t have any extra structure they don’t need.

How did we get to the definition of lists from the monoid laws, though? It doesn’t look anything like them. It would be nice if there was some systematic way to construct the corresponding free object given the laws of an algebra. Unfortunately, in normal Agda, this isn’t possible. Consider, for instance, if we added the commutativity law to the algebra: xy=yxx \bullet y = y \bullet x Not only is it not obvious how we’d write the corresponding free object, it’s actually not possible in normal Agda!

This kind of problem comes up a lot: we have a type, and we want it to obey just one more equation, but there is no inductive type which does so. Higher Inductive Types solve the problem in quite a straightforward way. So we want lists to satisfy another equation? Well, just add it to the definition!

module OddList where
 mutual
  data List {a} (A : Set a) : Set a where
    [] : List A
    _∷_ : A  List A  List A
    comm :  xs ys  xs ++ ys  ys ++ xs

  postulate _++_ : List A  List A  List A
Now, when we write a function that processes lists, Agda will check that the function behaves the same on xs ++ ys and ys ++ xs. As an example, here’s how you might define the free monoid as a HIT:
data FreeMonoid {a} (A : Set a) : Set a where
  [_] : A  FreeMonoid A
  _∙_ : FreeMonoid A  FreeMonoid A  FreeMonoid A
  ε : FreeMonoid A
  ∙ε :  x  x  ε  x
  ε∙ :  x  ε  x  x
  assoc :  x y z  (x  y)  z  x  (y  z)

It’s quite a satisfying definition, and very easy to see how we got to it from the monoid laws.

Now, when we write functions, we have to prove that those functions themselves also obey the monoid laws. For instance, here’s how we would take the length:
module Length where
  open import ProbabilityModule.Semirings.Nat
  open Semiring +-*-𝕊

  length : FreeMonoid A  
  length [ x ] = 1
  length (xs  ys) = length xs + length ys
  length ε = 0
  length (∙ε xs i) = +0 (length xs) i
  length (ε∙ xs i) = 0+ (length xs) i
  length (assoc xs ys zs i) = +-assoc (length xs) (length ys) (length zs) i

The first three clauses are the actual function: they deal with the three normal constructors of the type. The next three clauses prove that those previous clauses obey the equalities defined on the type.

With the preliminary stuff out of the way, let’s get on to the type I wanted to talk about:

Probability

First things first, let’s remember the classic definition of the probability monad:

newtype Prob a = Prob { runProb :: [(a, Rational)] }

Definitionally speaking, this doesn’t really represent what we’re talking about. For instance, the following two things express the same distribution, but have different representations:

Prob [(True, 1 / 4), (True, 1 / 4), (False, 1 / 2)]
Prob [(True , 1 / 2), (False, 1 / 2)]

So it’s the perfect candidate for an extra equality clause like we had above.

Second, in an effort to generalise, we won’t deal specifically with Rational, and instead we’ll use any semiring. After all of that, we get the following definition:

open Semiring rng

module Initial where
 infixr 5 _&_∷_
 data 𝒫 (A : Set a) : Set (a  s) where
   []  : 𝒫 A
   _&_∷_ : (p : R)  (x : A)  𝒫 A  𝒫 A
   dup :  p q x xs  p & x  q & x  xs  p + q & x  xs
   com :  p x q y xs  p & x  q & y  xs  q & y  p & x  xs
   del :  x xs  0# & x  xs  xs

The three extra conditions are pretty sensible: the first removes duplicates, the second makes things commutative, and the third removes impossible events.

Let’s get to writing some functions, then:

  : (A  R)  𝒫 A  R
  f [] = 0#
  f (p & x  xs) = p * f x +  f xs
  f (dup p q x xs i) = begin[ i ]
   p * f x + (q * f x +  f xs) ≡˘⟨ +-assoc (p * f x) (q * f x) ( f xs) 
   (p * f x + q * f x) +  f xs ≡˘⟨ cong (_+  f xs) (⟨+⟩* p q (f x))  
   (p + q) * f x +  f xs 
  f (com p x q y xs i) = begin[ i ]
   p * f x + (q * f y +  f xs) ≡˘⟨ +-assoc (p * f x) (q * f y) ( f xs) 
   p * f x + q * f y +  f xs   ≡⟨ cong (_+  f xs) (+-comm (p * f x) (q * f y)) 
   q * f y + p * f x +  f xs   ≡⟨ +-assoc (q * f y) (p * f x) ( f xs) 
   q * f y + (p * f x +  f xs) 
  f (del x xs i) = begin[ i ]
   0# * f x +  f xs ≡⟨ cong (_+  f xs) (0* (f x)) 
   0# +  f xs       ≡⟨ 0+ ( f xs) 
    f xs 

This is much more involved than the free monoid function, but the principle is the same: we first write the actual function (on the first three lines), and then we show that the function doesn’t care about the “rewrite rules” we have in the next three clauses.

Before going any further, we will have to amend the definition a little. The problem is that if we tried to prove something about any function on our 𝒫 type, we’d have to prove equalities between equalities as well. I’m sure that this is possible, but it’s very annoying, so I’m going to use a technique I saw in this repository. We add another rule to our type, stating that all equalities on the type are themselves equal. The new definition looks like this:

infixr 5 _&_∷_
data 𝒫 (A : Set a) : Set (a  s) where
  []  : 𝒫 A
  _&_∷_ : (p : R)  (x : A)  𝒫 A  𝒫 A
  dup :  p q x xs  p & x  q & x  xs  p + q & x  xs
  com :  p x q y xs  p & x  q & y  xs  q & y  p & x  xs
  del :  x xs  0# & x  xs  xs
  trunc : isSet (𝒫 A)

Eliminators

Unfortunately, after adding that case we have to deal with it explicitly in every pattern-match on 𝒫. We can get around it by writing an eliminator for the type which deals with it itself. Eliminators are often irritating to work with, though: we give up the nice pattern-matching syntax we get when we program directly. It’s a bit like having to rely on church encoding everywhere.

However, we can get back some pattern-like syntax if we use copatterns. Here’s an example of what I mean, for folds on lists:

module ListElim where
 open NormalList
 open import ProbabilityModule.Semirings.Nat
 open Semiring +-*-𝕊 renaming (_+_ to _ℕ+_)

 record [_↦_] (A : Set a) (B : Set b) : Set (a  b) where
   field
     [_][] : B
     [_]_∷_ : A  B  B
   [_]↓ : List A  B
   [ [] ]↓ = [_][]
   [ x  xs ]↓ = [_]_∷_ x [ xs ]↓
 open [_↦_]
 
 sum-alg : [    ]
 [ sum-alg ][] = 0
 [ sum-alg ] x  xs = x ℕ+ xs
 
 sum : List   
 sum = [ sum-alg ]↓

For the probability monad, there’s an eliminator for the whole thing, and eliminator for propositional proofs, and a normal eliminator for folding. Their definitions are quite long, but mechanical.

Eliminator Definitions
record ⟅_↝_⟆ {a } (A : Set a) (P : 𝒫 A  Set ) : Set (a    s) where
  constructor elim
  field
    ⟅_⟆-set :  {xs}  isSet (P xs)
    ⟅_⟆[] : P []
    ⟅_⟆_&_∷_ :  p x xs  P xs  P (p & x  xs)
  private z = ⟅_⟆[]; f = ⟅_⟆_&_∷_
  field
    ⟅_⟆-dup : (∀ p q x xs pxs  PathP  i  P (dup p q x xs i))
              (f p x (q & x  xs) (f q x xs pxs)) (f (p + q) x xs pxs))
    ⟅_⟆-com : (∀ p x q y xs pxs  PathP  i  P (com p x q y xs i))
              (f p x (q & y  xs) (f q y xs pxs)) (f q y (p & x  xs) (f p x xs pxs)))
    ⟅_⟆-del : (∀ x xs pxs  PathP  i  P (del x xs i))
              (f 0# x xs pxs) pxs)
  ⟅_⟆⇓ : (xs : 𝒫 A)  P xs
   [] ⟆⇓ = z
   p & x  xs ⟆⇓ = f p x xs  xs ⟆⇓
   dup p q x xs i ⟆⇓ = ⟅_⟆-dup p q x xs  xs ⟆⇓ i
   com p x q y xs i ⟆⇓ = ⟅_⟆-com p x q y xs  xs ⟆⇓ i
   del x xs i ⟆⇓ = ⟅_⟆-del x xs  xs ⟆⇓ i
   trunc xs ys p q i j ⟆⇓ =
    elimSquash₀  xs  ⟅_⟆-set {xs}) (trunc xs ys p q)  xs ⟆⇓  ys ⟆⇓ (cong ⟅_⟆⇓ p) (cong ⟅_⟆⇓ q) i j

open ⟅_↝_⟆ public
elim-syntax :  {a }  (A : Set a)  (𝒫 A  Set )  Set (a    s)
elim-syntax = ⟅_↝_⟆

syntax elim-syntax A  xs  Pxs) = [ xs ∈𝒫 A  Pxs ]

record ⟦_⇒_⟧ {a } (A : Set a) (P : 𝒫 A  Set ) : Set (a    s) where
  constructor elim-prop
  field
    ⟦_⟧-prop :  {xs}  isProp (P xs)
    ⟦_⟧[] : P []
    ⟦_⟧_&_∷_⟨_⟩ :  p x xs  P xs  P (p & x  xs)
  private z = ⟦_⟧[]; f = ⟦_⟧_&_∷_⟨_⟩
  ⟦_⟧⇑ = elim
          (isProp→isSet ⟦_⟧-prop)
          z f
           p q x xs pxs 
            toPathP (⟦_⟧-prop (transp  i  P (dup p q x xs i))
            i0
            (f p x (q & x  xs) (f q x xs pxs))) (f (p + q) x xs pxs) ))
           p x q y xs pxs  toPathP (⟦_⟧-prop (transp  i  P (com p x q y xs i)) i0
            (f p x (q & y  xs) (f q y xs pxs))) (f q y (p & x  xs) (f p x xs pxs))))
            λ x xs pxs  toPathP (⟦_⟧-prop (transp  i  P (del x xs i)) i0 ((f 0# x xs pxs))) pxs)
  ⟦_⟧⇓ =  ⟦_⟧⇑ ⟆⇓

open ⟦_⇒_⟧ public
elim-prop-syntax :  {a }  (A : Set a)  (𝒫 A  Set )  Set (a    s)
elim-prop-syntax = ⟦_⇒_⟧

syntax elim-prop-syntax A  xs  Pxs) =  xs ∈𝒫 A  Pxs 

record [_↦_] {a b} (A : Set a) (B : Set b) : Set (a  b  s) where
  constructor rec
  field
    [_]-set  : isSet B
    [_]_&_∷_ : R  A  B  B
    [_][]    : B
  private f = [_]_&_∷_; z = [_][]
  field
    [_]-dup  :  p q x xs  f p x (f q x xs)  f (p + q) x xs
    [_]-com :  p x q y xs  f p x (f q y xs)  f q y (f p x xs)
    [_]-del :  x xs  f 0# x xs  xs
  [_]⇑ = elim
            [_]-set
            z
             p x _ xs  f p x xs)
             p q x xs  [_]-dup p q x)
             p x q y xs  [_]-com p x q y)
             x xs  [_]-del x)
  [_]↓ =  [_]⇑ ⟆⇓
open [_↦_] public

Here’s one in action, to define map:

map : (A  B)  𝒫 A  𝒫 B
map = λ f  [ map′ f ]↓
  module Map where
  map′ : (A  B)  [ A  𝒫 B ]
  [ map′ f ] p & x  xs = p & f x  xs
  [ map′ f ][] = []
  [ map′ f ]-set = trunc
  [ map′ f ]-dup p q x xs = dup p q (f x) xs
  [ map′ f ]-com p x q y xs = com p (f x) q (f y) xs
  [ map′ f ]-del x xs = del (f x) xs

And here’s how we’d define union, and then prove that it’s associative:

infixr 5 _∪_
_∪_ : 𝒫 A  𝒫 A  𝒫 A
_∪_ = λ xs ys  [ union ys ]↓ xs
  module Union where
  union : 𝒫 A  [ A  𝒫 A ]
  [ union ys ]-set = trunc
  [ union ys ] p & x  xs = p & x  xs
  [ union ys ][] = ys
  [ union ys ]-dup = dup
  [ union ys ]-com = com
  [ union ys ]-del = del

∪-assoc : (xs ys zs : 𝒫 A)  xs  (ys  zs)  (xs  ys)  zs
∪-assoc = λ xs ys zs   ∪-assoc′ ys zs ⟧⇓ xs
  module UAssoc where
  ∪-assoc′ :  ys zs   xs ∈𝒫 A  xs  (ys  zs)  (xs  ys)  zs 
   ∪-assoc′ ys zs ⟧-prop = trunc _ _
   ∪-assoc′ ys zs ⟧[] = refl
   ∪-assoc′ ys zs  p & x  xs  P  = cong (p & x ∷_) P

There’s a lot more stuff here that I won’t bore you with.

Boring Stuff
infixl 7 _⋊_
_⋊_ : R  𝒫 A  𝒫 A
_⋊_ = λ p  [ p ⋊′ ]↓
  module Cond where
  _⋊′ : R  [ A  𝒫 A ]
  [ p ⋊′ ]-set = trunc
  [ p ⋊′ ][] = []
  [ p ⋊′ ] q & x  xs = p * q & x  xs
  [ p ⋊′ ]-com q x r y xs = com (p * q) x (p * r) y xs
  [ p ⋊′ ]-dup q r x xs =
    p * q & x  p * r & x  xs ≡⟨ dup (p * q) (p * r) x xs 
    p * q + p * r & x  xs     ≡˘⟨ cong (_& x  xs) (*⟨+⟩ p q r) 
    p * (q + r) & x  xs       
  [ p ⋊′ ]-del x xs =
    p * 0# & x  xs ≡⟨ cong (_& x  xs) (*0 p) 
    0# & x  xs     ≡⟨ del x xs 
    xs              

 : (A  R)  𝒫 A  R
 = λ f  [ ∫′ f ]↓
  module Expect where
  ∫′ : (A  R)  [ A  R ]
  [ ∫′ f ]-set = sIsSet
  [ ∫′ f ] p & x  xs = p * f x + xs
  [ ∫′ f ][] = 0#
  [ ∫′ f ]-dup p q x xs =
    p * f x + (q * f x + xs) ≡˘⟨ +-assoc (p * f x) (q * f x) xs 
    (p * f x + q * f x) + xs ≡˘⟨ cong (_+ xs) (⟨+⟩* p q (f x)) 
    (p + q) * f x + xs 
  [ ∫′ f ]-com p x q y xs =
    p * f x + (q * f y + xs) ≡˘⟨ +-assoc (p * f x) (q * f y) (xs) 
    p * f x + q * f y + xs   ≡⟨ cong (_+ xs) (+-comm (p * f x) (q * f y)) 
    q * f y + p * f x + xs   ≡⟨ +-assoc (q * f y) (p * f x) (xs) 
    q * f y + (p * f x + xs) 
  [ ∫′ f ]-del x xs =
    0# * f x + xs ≡⟨ cong (_+ xs) (0* (f x)) 
    0# + xs       ≡⟨ 0+ (xs) 
    xs            

syntax   x  e) =  e 𝑑 x

pure : A  𝒫 A
pure x = 1# & x  []

∪-cons :  p (x : A) xs ys  xs  p & x  ys  p & x  xs  ys
∪-cons = λ p x xs ys   ∪-cons′ p x ys ⟧⇓ xs
  module UCons where
  ∪-cons′ :  p x ys   xs ∈𝒫 A  xs  p & x  ys  p & x  xs  ys 
   ∪-cons′ p x ys ⟧-prop = trunc _ _
   ∪-cons′ p x ys ⟧[] = refl
   ∪-cons′ p x ys  r & y  xs  P  = cong (r & y ∷_) P ; com r y p x (xs  ys)

⋊-distribʳ :  p q  (xs : 𝒫 A)  p  xs  q  xs  (p + q)  xs
⋊-distribʳ = λ p q   ⋊-distribʳ′ p q ⟧⇓
  module JDistrib where
  ⋊-distribʳ′ :  p q   xs ∈𝒫 A  p  xs  q  xs  (p + q)  xs 
   ⋊-distribʳ′ p q ⟧-prop = trunc _ _
   ⋊-distribʳ′ p q ⟧[] = refl
   ⋊-distribʳ′ p q  r & x  xs  P  =
    p  (r & x  xs)  q  (r & x  xs)   ≡⟨ ∪-cons (q * r) x (p  (r & x  xs)) (q  xs)  
    q * r & x  p  (r & x  xs)  q  xs ≡⟨ cong (_∪ q  xs) (dup (q * r) (p * r) x (p  xs)) 
    q * r + p * r & x  p  xs  q  xs   ≡˘⟨ cong (_& x  (p  xs  q  xs)) (⟨+⟩* q p r) 
    (q + p) * r & x  p  xs  q  xs     ≡⟨ cong ((q + p) * r & x ∷_) P 
    (q + p) * r & x  (p + q)  xs        ≡⟨ cong  pq  pq * r & x  (p + q)  xs) (+-comm q p) 
    (p + q) * r & x  (p + q)  xs        ≡⟨⟩
    _⋊_ (p + q) (r & x  xs) 

⋊-distribˡ :  p  (xs ys : 𝒫 A)  p  xs  p  ys  p  (xs  ys)
⋊-distribˡ = λ p xs ys   ⋊-distribˡ′ p ys ⟧⇓ xs
  module JDistribL where
  ⋊-distribˡ′ :  p ys   xs ∈𝒫 A  p  xs  p  ys  p  (xs  ys) 
   ⋊-distribˡ′ p ys ⟧-prop = trunc _ _
   ⋊-distribˡ′ p ys ⟧[] = refl
   ⋊-distribˡ′ p ys  q & x  xs  P  =
    p  (q & x  xs)  p  ys ≡⟨⟩
    p * q & x  p  xs  p  ys ≡⟨ cong (p * q & x ∷_) P 
    p * q & x  p  (xs  ys) ≡⟨⟩
    p  ((q & x  xs)  ys) 


∪-idʳ : (xs : 𝒫 A)  xs  []  xs
∪-idʳ =  ∪-idʳ′ ⟧⇓
  module UIdR where
  ∪-idʳ′ :  xs ∈𝒫 A  xs  []  xs 
   ∪-idʳ′ ⟧-prop = trunc _ _
   ∪-idʳ′ ⟧[] = refl
   ∪-idʳ′  p & x  xs  P  = cong (p & x ∷_) P

∪-comm : (xs ys : 𝒫 A)  xs  ys  ys  xs
∪-comm = λ xs ys   ∪-comm′ ys ⟧⇓ xs
  module UComm where
  ∪-comm′ :  ys   xs ∈𝒫 A  xs  ys  ys  xs 
   ∪-comm′ ys ⟧-prop = trunc _ _
   ∪-comm′ ys ⟧[] = sym (∪-idʳ ys)
   ∪-comm′ ys  p & x  xs  P  = cong (p & x ∷_) P ; sym (∪-cons p x ys xs)

0⋊ : (xs : 𝒫 A)  0#  xs  []
0⋊ =  0⋊′ ⟧⇓
  module ZeroJ where
  0⋊′ :  xs ∈𝒫 A  0#  xs  [] 
   0⋊′ ⟧-prop = trunc _ _
   0⋊′ ⟧[] = refl
   0⋊′  p & x  xs  P  =
    0#  (p & x  xs)    ≡⟨⟩
    0# * p & x  0#  xs ≡⟨ cong (_& x  0#  xs) (0* p) 
    0# & x  0#  xs     ≡⟨ del x (0#  xs) 
    0#  xs              ≡⟨ P 
    [] 

However, I can demonstrate the monadic bind:

_>>=_ : 𝒫 A  (A  𝒫 B)  𝒫 B
xs >>= f = [ f =<< ]↓ xs
  module Bind where
  _=<< : (A  𝒫 B)  [ A  𝒫 B ]
  [ f =<< ] p & x  xs = p  (f x)  xs
  [ f =<< ][] = []
  [ f =<< ]-set = trunc
  [ f =<< ]-del x xs = cong (_∪ xs) (0⋊ (f x))
  [ f =<< ]-dup p q x xs =
    p  (f x)  q  (f x)  xs   ≡⟨ ∪-assoc (p  f x) (q  f x) xs 
    (p  (f x)  q  (f x))  xs ≡⟨ cong (_∪ xs) (⋊-distribʳ p q (f x) ) 
    _⋊_ (p + q) (f x)  xs 
  [ f =<< ]-com p x q y xs =
    p  (f x)  q  (f y)  xs   ≡⟨ ∪-assoc (p  f x) (q  f y) xs 
    (p  (f x)  q  (f y))  xs ≡⟨ cong (_∪ xs) (∪-comm (p  f x) (q  f y)) 
    (q  (f y)  p  (f x))  xs ≡˘⟨ ∪-assoc (q  f y) (p  f x) xs 
    q  (f y)  p  (f x)  xs 

And we can prove the monad laws, also:

Proofs of Monad Laws
1⋊ : (xs : 𝒫 A)  1#  xs  xs
1⋊ =  1⋊′ ⟧⇓
  module OneJoin where
  1⋊′ :  xs ∈𝒫 A  1#  xs  xs 
   1⋊′ ⟧-prop = trunc _ _
   1⋊′ ⟧[] = refl
   1⋊′  p & x  xs  P  =
    1#  (p & x  xs) ≡⟨⟩
    1# * p & x  1#  xs ≡⟨ cong (_& x  1#  xs) (1* p) 
    p & x  1#  xs ≡⟨ cong (p & x ∷_) P 
    p & x  xs 

>>=-distrib : (xs ys : 𝒫 A) (g : A  𝒫 B)  (xs  ys) >>= g  (xs >>= g)  (ys >>= g)
>>=-distrib = λ xs ys g   >>=-distrib′ ys g ⟧⇓ xs
  module BindDistrib where
  >>=-distrib′ : (ys : 𝒫 A) (g : A  𝒫 B)   xs ∈𝒫 A  ((xs  ys) >>= g)  (xs >>= g)  (ys >>= g) 
   >>=-distrib′ ys g ⟧-prop = trunc _ _
   >>=-distrib′ ys g ⟧[] = refl
   >>=-distrib′ ys g  p & x  xs  P  =
    (((p & x  xs)  ys) >>= g) ≡⟨⟩
    (p & x  xs  ys) >>= g ≡⟨⟩
    p  g x  ((xs  ys) >>= g) ≡⟨ cong (p  g x ∪_) P 
    p  g x  ((xs >>= g)  (ys >>= g)) ≡⟨ ∪-assoc (p  g x) (xs >>= g) (ys >>= g) 
    (p  g x  (xs >>= g))  (ys >>= g) ≡⟨⟩
    ((p & x  xs) >>= g)  (ys >>= g) 

*-assoc-⋊ :  p q (xs : 𝒫 A)  (p * q)  xs  p  (q  xs)
*-assoc-⋊ = λ p q   *-assoc-⋊′ p q ⟧⇓
  module MAssocJ where
  *-assoc-⋊′ :  p q   xs ∈𝒫 A  (p * q)  xs  p  (q  xs) 
   *-assoc-⋊′ p q ⟧-prop = trunc _ _
   *-assoc-⋊′ p q ⟧[] = refl
   *-assoc-⋊′ p q  r & x  xs  P  =
    p * q  (r & x  xs) ≡⟨⟩
    p * q * r & x  (p * q  xs) ≡⟨ cong (_& x  (p * q  xs)) (*-assoc p q r) 
    p * (q * r) & x  (p * q  xs) ≡⟨ cong (p * (q * r) & x ∷_) P 
    p * (q * r) & x  (p  (q  xs)) ≡⟨⟩
    p  (q  (r & x  xs)) 

⋊-assoc->>= :  p (xs : 𝒫 A) (f : A  𝒫 B)  (p  xs) >>= f  p  (xs >>= f)
⋊-assoc->>= = λ p xs f   ⋊-assoc->>=′ p f ⟧⇓ xs
  module JDistribB where
  ⋊-assoc->>=′ :  p (f : A  𝒫 B)   xs ∈𝒫 A  (p  xs) >>= f  p  (xs >>= f) 
   ⋊-assoc->>=′ p f ⟧-prop = trunc _ _
   ⋊-assoc->>=′ p f ⟧[] = refl
   ⋊-assoc->>=′ p f  q & x  xs  P  =
    (p  (q & x  xs)) >>= f ≡⟨⟩
    (p * q & x  p  xs) >>= f ≡⟨⟩
    ((p * q)  f x)  ((p  xs) >>= f) ≡⟨ cong (((p * q)  f x) ∪_) P 
    ((p * q)  f x)  (p  (xs >>= f)) ≡⟨ cong (_∪ (p  (xs >>= f))) (*-assoc-⋊ p q (f x)) 
    (p  (q  f x))  (p  (xs >>= f)) ≡⟨ ⋊-distribˡ p (q  f x) (xs >>= f) 
    p  ((q & x  xs) >>= f) 

>>=-idˡ : (x : A)  (f : A  𝒫 B)
       (pure x >>= f)  f x
>>=-idˡ x f =
  pure x >>= f ≡⟨⟩
  (1# & x  []) >>= f ≡⟨⟩
  1#  f x  [] >>= f ≡⟨⟩
  1#  f x  [] ≡⟨ ∪-idʳ (1#  f x) 
  1#  f x ≡⟨ 1⋊ (f x) 
  f x 

>>=-idʳ : (xs : 𝒫 A)  xs >>= pure  xs
>>=-idʳ =  >>=-idʳ′ ⟧⇓
  module Law1 where
  >>=-idʳ′ :  xs ∈𝒫 A  xs >>= pure  xs 
   >>=-idʳ′ ⟧-prop = trunc _ _
   >>=-idʳ′ ⟧[] = refl
   >>=-idʳ′  p & x  xs  P  =
    ((p & x  xs) >>= pure) ≡⟨⟩
    p  (pure x)  (xs >>= pure) ≡⟨⟩
    p  (1# & x  [])  (xs >>= pure) ≡⟨⟩
    p * 1# & x  []  (xs >>= pure) ≡⟨⟩
    p * 1# & x  (xs >>= pure) ≡⟨ cong (_& x  (xs >>= pure)) (*1 p) 
    p & x  xs >>= pure ≡⟨ cong (p & x ∷_) P 
    p & x  xs 

>>=-assoc : (xs : 𝒫 A)  (f : A  𝒫 B)  (g : B  𝒫 C)
       ((xs >>= f) >>= g)  xs >>=  x  f x >>= g)
>>=-assoc = λ xs f g   >>=-assoc′ f g ⟧⇓ xs
  module Law3 where
  >>=-assoc′ : (f : A  𝒫 B)  (g : B  𝒫 C)   xs ∈𝒫 A  ((xs >>= f) >>= g)  xs >>=  x  f x >>= g) 
   >>=-assoc′ f g ⟧-prop = trunc _ _
   >>=-assoc′ f g ⟧[] = refl
   >>=-assoc′ f g  p & x  xs  P  =
    (((p & x  xs) >>= f) >>= g) ≡⟨⟩
    ((p  f x  (xs >>= f)) >>= g) ≡⟨ >>=-distrib (p  f x) (xs >>= f) g 
    ((p  f x) >>= g)  ((xs >>= f) >>= g) ≡⟨ cong ((p  f x) >>= g ∪_) P 
    ((p  f x) >>= g)  (xs >>=  y  f y >>= g)) ≡⟨ cong (_∪ (xs >>=  y  f y >>= g))) (⋊-assoc->>= p (f x) g) 
    p  (f x >>= g)  (xs >>=  y  f y >>= g)) ≡⟨⟩
    ((p & x  xs) >>=  y  f y >>= g)) 

Conclusion

I’ve really enjoyed working with cubical Agda so far, and the proofs above were a pleasure to write. I think I can use the above definition to get a workable differential privacy monad, also.

Anyway, all the code is available here.