fix aggregated posterior computation as described in #3188

scverse · Feb 19, 2025 · 09947b6 · 09947b6
1 parent 126e0d3
commit 09947b6
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/src/scvi/external/mrvi/_model.py b/src/scvi/external/mrvi/_model.py
@@ -789,11 +789,11 @@ def get_aggregated_posterior(
             qu_locs.append(outputs["qu"].loc)
             qu_scales.append(outputs["qu"].scale)
 
-        qu_loc = jnp.concatenate(qu_locs, axis=0).T
-        qu_scale = jnp.concatenate(qu_scales, axis=0).T
+        qu_loc = jnp.concatenate(qu_locs, axis=0)
+        qu_scale = jnp.concatenate(qu_scales, axis=0)
         return MixtureSameFamily(
-            Categorical(probs=jnp.ones(qu_loc.shape[1]) / qu_loc.shape[1]),
-            Normal(qu_loc, qu_scale),
+            Categorical(probs=jnp.ones(qu_loc.shape[0]) / qu_loc.shape[0]),
+            Normal(qu_loc, qu_scale).to_event(1),
         )
 
     def differential_abundance(