Document how particles should be represented in SMC

blackjax/smc/base.py

@@ -20,7 +20,21 @@
 
 
 class SMCState(NamedTuple):
-    """State of the SMC sampler"""
+    """State of the SMC sampler
+    Particles must be a ArrayTree, each leave represents
+    a variable from the posterior, being an array of size (n_particles, the variable's dimensionality).
+    Examples (three particles):
+        - Single univariate posterior:
+            [ Array([[1.], [1.2], [3.4]]) ]
+        - Single bivariate  posterior:
+            [Array([[1,2], [3,4], [5,6]])]
+        - Two variables, each univariate:
+            [ Array([[1.], [1.2], [3.4]]),
+            Array([[50.], [51], [55]]) ]
+        - Two variables, first one bivariate, second one 4-variate:
+            [ Array([[1., 2.], [1.2, 0.5], [3.4, 50]]),
+            Array([[50., 51., 52., 51], [51., 52., 52. ,54.], [55., 60, 60, 70]])]
+    """
 
     particles: ArrayTree
     weights: Array