Correct typo and make error message more specific.

astro-informatics · Aug 14, 2024 · 631cc00 · 631cc00
1 parent 58dd63f
commit 631cc00
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/harmonic/model.py b/harmonic/model.py
@@ -112,7 +112,7 @@ def __init__(
 
         if transformation != None and log_J_det == None:
             raise ValueError(
-                "Please specify function to calculate Jacobian determinant of data transformation."
+                "Please specify function to calculate log of absolute value of the determinant Jacobian determinant of data transformation."
             )
 
         self.ndim = ndim_in
@@ -126,8 +126,8 @@ def __init__(
         self.temperature = temperature
         self.flow = None
         if not transformation is None:
-            self.transformation_vmap = jax.vmap(transformation, in_axes=0)
-            self.J_det_vmap = jax.vmap(log_J_det, in_axes=0)
+            self.transformation_vmap = jax.vmap(transformation, in_axes=(0, None))
+            self.log_J_det_vmap = jax.vmap(log_J_det, in_axes=(0, None))
         else:
             self.transformation_vmap = transformation