From fba5f791de0cbdd61898c1aced8a82c0463a2706 Mon Sep 17 00:00:00 2001
From: Thomas Ng <thomas.ng@link.cuhk.edu.hk>
Date: Thu, 4 Jul 2024 10:37:43 +0800
Subject: [PATCH] Update prior.py

---
 src/jimgw/prior.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/src/jimgw/prior.py b/src/jimgw/prior.py
index bb43b9b2..b95eb7bc 100644
--- 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.