Update prior.py

kazewong · Jul 4, 2024 · fba5f79 · fba5f79
1 parent e8e645c
commit fba5f79
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/src/jimgw/prior.py b/src/jimgw/prior.py
@@ -224,6 +224,7 @@ def log_prob(self, x: dict[str, Float]) -> Float:
         return jnp.log(jnp.exp(-variable) / (1 + jnp.exp(-variable)) ** 2)
 
 
+@jaxtyped(typechecker=typechecker)
 class Sphere(Prior):
     """
     A prior on a sphere represented by Cartesian coordinates.
@@ -395,6 +396,8 @@ def log_prob(self, x: dict[str, Float]) -> Float:
         )
         return log_p
 
+
+@jaxtyped(typechecker=typechecker)
 class EarthFrame(Prior):
     """
     Prior distribution for sky location in Earth frame.