chore: lax.sqaure replaced with jnp.sqaure

pyro-ppl · Jul 6, 2024 · 1192ca7 · 1192ca7
1 parent d8eb1e5
commit 1192ca7
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/numpyro/distributions/continuous.py b/numpyro/distributions/continuous.py
@@ -3009,7 +3009,7 @@ def mean(self):
     @lazy_property
     def variance(self):
         moment2 = self._kth_moment(2)
-        return moment2 - lax.square(self.mean)
+        return moment2 - jnp.square(self.mean)
 
 
 class LowerTruncatedPowerLaw(Distribution):
@@ -3100,7 +3100,7 @@ def mean(self):
 
     @lazy_property
     def variance(self):
-        return self._kth_moment(2) - lax.square(self.mean)
+        return self._kth_moment(2) - jnp.square(self.mean)
 
     def entropy(self):
         # The simplified expression for the entorpy is,