Updated jim.py

xuyuon · Sep 4, 2024 · 1d00fd6 · 1d00fd6
1 parent 21f1b38
commit 1d00fd6
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 3 deletions.
diff --git a/example/GW150914_Pv2_reparam.py b/example/GW150914_Pv2_reparam.py
@@ -150,7 +150,7 @@
 
 jim.sample(jax.random.PRNGKey(42))
 jim.get_samples()
-jim.print_summary()
+# jim.print_summary()
 
 ###########################################
 ########## Visualize the Data #############
@@ -164,7 +164,7 @@
 if jim.sample_transforms:
     transformed_chain = jim.add_name(production_chain)
     for transform in reversed(jim.sample_transforms):
-        transformed_chain = transform.backward(transformed_chain)
+        transformed_chain = jax.vmap(transform.backward)(transformed_chain)
 result = transformed_chain
 labels = list(transformed_chain.keys())
 

diff --git a/src/jimgw/jim.py b/src/jimgw/jim.py
@@ -224,8 +224,12 @@ def get_samples(self, training: bool = False) -> dict:
 
         chains = chains.transpose(2, 0, 1)
         chains = self.add_name(chains)
+        # flatten the shape
+        chains = {key: value.reshape(-1) for key, value in chains.items()}
         for sample_transform in reversed(self.sample_transforms):
-            chains = sample_transform.backward(chains)
+            chains = jax.vmap(sample_transform.backward)(chains)
+        # restore the shape
+        chains = {key: value.reshape(-1, self.sampler.n_chains).T for key, value in chains.items()}
         return chains
 
     def plot(self):