SteinVI: Allow nested params

numpyro/contrib/einstein/steinvi.py

@@ -111,16 +111,22 @@ def _param_size(self, param):
             return sum(map(self._param_size, param))
         return param.size
 
-    def _calc_particle_info(self, uparams, num_particles):
+    def _calc_particle_info(self, uparams, num_particles, start_index=0):
         uparam_keys = list(uparams.keys())
         uparam_keys.sort()
-        start_index = 0
         res = {}
+        end_index = start_index
         for k in uparam_keys:
-            end_index = start_index + self._param_size(uparams[k]) // num_particles
-            res[k] = (start_index, end_index)
+            if isinstance(uparams[k], dict):
+                res_sub, end_index = self._calc_particle_info(
+                    uparams[k], num_particles, start_index
+                )
+                res[k] = res_sub
+            else:
+                end_index = start_index + self._param_size(uparams[k]) // num_particles
+                res[k] = (start_index, end_index)
             start_index = end_index
-        return res
+        return res, end_index
 
     def _find_init_params(self, particle_seed, inner_guide, inner_guide_trace):
         def extract_info(site):
@@ -180,7 +186,7 @@ def _svgd_loss_and_grads(self, rng_key, unconstr_params, *args, **kwargs):
         stein_particles, unravel_pytree, unravel_pytree_batched = batch_ravel_pytree(
             stein_uparams, nbatch_dims=1
         )
-        particle_info = self._calc_particle_info(
+        particle_info, _ = self._calc_particle_info(
             stein_uparams, stein_particles.shape[0]
         )
 

test/contrib/einstein/test_stein.py

@@ -382,12 +382,35 @@ def test_calc_particle_info(num_params, num_particles):
     expected_pinfo = dict(zip(string.ascii_lowercase[:num_params], expected_start_end))
 
     stein = SteinVI(id, id, Adam(1.0), Trace_ELBO(), RBFKernel())
-    pinfo = stein._calc_particle_info(uparams, num_particles)
+    pinfo, _ = stein._calc_particle_info(uparams, num_particles)
 
     for k in pinfo.keys():
         assert pinfo[k] == expected_pinfo[k], f"Failed for seed {seed}"
 
 
+def test_calc_particle_info_nested():
+    num_params = 3
+    num_particles = 10
+    seed = random.PRNGKey(42)
+    sizes = Poisson(5).sample(seed, (100, nrandom.randint(1, 10))) + 1
+    uparam = tuple(np.empty(tuple(size)) for size in sizes)
+    uparams = {
+        string.ascii_lowercase[i]: {
+            string.ascii_lowercase[j]: uparam for j in range(num_params)
+        }
+        for i in range(num_params)
+    }
+
+    stein = SteinVI(id, id, Adam(1.0), Trace_ELBO(), RBFKernel())
+    pinfo, _ = stein._calc_particle_info(uparams, num_particles)
+    start = 0
+    tot_size = sum(map(lambda size: size.prod(), sizes)) // num_particles
+    for val in pinfo.values():
+        for v in val.values():
+            assert v == (start, start + tot_size)
+            start += tot_size
+
+
 ########################################
 # Stein Kernels
 ########################################