bayinx 0.2.28__py3-none-any.whl → 0.2.29__py3-none-any.whl

bayinx/dists/gamma2.py

@@ -0,0 +1,39 @@
+import jax.lax as lax
+from jax.scipy.special import gammaln
+from jaxtyping import Array, ArrayLike, Float, Real
+
+
+def prob(
+    x: Real[ArrayLike, "..."], mu: Real[ArrayLike, "..."], nu: Real[ArrayLike, "..."]
+) -> Float[Array, "..."]:
+    """
+    The probability density function (PDF) for a (mean-precision parameterized) Gamma distribution.
+
+    # Parameters
+    - `x`:      Value(s) at which to evaluate the PDF.
+    - `mu`:     The mean.
+    - `nu`:     The positive inverse dispersion.
+
+    # Returns
+    The PDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `mu`, and `nu`.
+    """
+
+    return lax.exp(logprob(x, mu, nu))
+
+
+def logprob(
+    x: Real[ArrayLike, "..."], mu: Real[ArrayLike, "..."], nu: Real[ArrayLike, "..."]
+) -> Float[Array, "..."]:
+    """
+    The log-transformed probability density function (log PDF) for a (mean-precision parameterized) Gamma distribution.
+
+    # Parameters
+    - `x`:      Value(s) at which to evaluate the log PDF.
+    - `mu`:     The mean/location.
+    - `nu`:     The positive inverse dispersion.
+
+    # Returns
+    The log PDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `mu`, and `nu`.
+    """
+
+    return - gammaln(nu) + nu * (lax.log(nu) - lax.log(mu)) + (nu - 1.0) * lax.log(x) - (x * nu / mu) # pyright: ignore

bayinx/dists/normal.py

@@ -1,4 +1,4 @@
-import jax.lax as _lax
+import jax.lax as lax
 from jaxtyping import Array, ArrayLike, Float, Real
 
 __PI = 3.141592653589793
@@ -12,15 +12,15 @@ def prob(
 
     # Parameters
     - `x`:      Value(s) at which to evaluate the PDF.
-    - `mu`:     The mean/location parameter(s).
-    - `sigma`:  The non-negative standard deviation parameter(s).
+    - `mu`:     The mean/location.
+    - `sigma`:  The positive standard deviation.
 
     # Returns
     The PDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `mu`, and `sigma`.
     """
 
-    return _lax.exp(-0.5 * _lax.square((x - mu) / sigma)) / (  # pyright: ignore
-        sigma * _lax.sqrt(2.0 * __PI)
+    return lax.exp(-0.5 * lax.square((x - mu) / sigma)) / (  # pyright: ignore
+        sigma * lax.sqrt(2.0 * __PI)
     )
 
 
@@ -36,11 +36,11 @@ def logprob(
     - `sigma`:  The non-negative standard deviation parameter(s).
 
     # Returns
-    The log of the PDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `mu`, and `sigma`.
+    The log PDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `mu`, and `sigma`.
     """
 
-    return -_lax.log(sigma * _lax.sqrt(2.0 * __PI)) - 0.5 * _lax.square(
-        (x - mu) / sigma # pyright: ignore
+    return -lax.log(sigma * lax.sqrt(2.0 * __PI)) - 0.5 * lax.square(
+        (x - mu) / sigma  # pyright: ignore
     )
 
 
@@ -53,13 +53,13 @@ def uprob(
     # Parameters
     - `x`:      Value(s) at which to evaluate the uPDF.
     - `mu`:     The mean/location parameter(s).
-    - `sigma`:  The non-negative standard deviation parameter(s).
+    - `sigma`:  The positive standard deviation parameter(s).
 
     # Returns
     The uPDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `mu`, and `sigma`.
     """
 
-    return _lax.exp(-0.5 * _lax.square((x - mu) / sigma)) / sigma  # pyright: ignore
+    return lax.exp(-0.5 * lax.square((x - mu) / sigma)) / sigma  # pyright: ignore
 
 
 def ulogprob(
@@ -77,4 +77,4 @@ def ulogprob(
     The log uPDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `mu`, and `sigma`.
     """
 
-    return -_lax.log(sigma) - 0.5 * _lax.square((x - mu) / sigma)  # pyright: ignore
+    return -lax.log(sigma) - 0.5 * lax.square((x - mu) / sigma)  # pyright: ignore

bayinx/mhx/vi/flows/fullaffine.py

@@ -30,19 +30,19 @@ class FullAffine(Flow):
             "scale": jnp.zeros((dim, dim)),
         }
 
-        self.constraints = {"scale": lambda m: jnp.tril(m)}
+        if dim == 1:
+            self.constraints = {}
+        else:
 
-    @eqx.filter_jit
-    def transform_pars(self):
-        params = self.constrain_pars()
-
-        # Extract diagonal and apply exponential
-        diag: Array = jnp.exp(jnp.diag(params["scale"]))
+            @eqx.filter_jit
+            def constrain_scale(scale: Array):
+                # Extract diagonal and apply exponential
+                diag: Array = jnp.exp(jnp.diag(scale))
 
-        # Fill diagonal
-        params["scale"] = jnp.fill_diagonal(params["scale"], diag, inplace=False)
+                # Return matrix with modified diagonal
+                return jnp.fill_diagonal(scale, diag, inplace=False)
 
-        return params
+            self.constraints = {"scale": constrain_scale}
 
     @eqx.filter_jit
     def forward(self, draws: Array) -> Array:

{bayinx-0.2.28.dist-info → bayinx-0.2.29.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bayinx
-Version: 0.2.28
+Version: 0.2.29
 Summary: Bayesian Inference with JAX
 Requires-Python: >=3.12
 Requires-Dist: equinox>=0.11.12

{bayinx-0.2.28.dist-info → bayinx-0.2.29.dist-info}/RECORD

@@ -7,10 +7,8 @@ bayinx/core/model.py,sha256=Z_HaFr0_-keMjG5tg3xxP3hGML7aDFIcCI8Y5dGrtM4,2145
 bayinx/core/variational.py,sha256=W0747jfVJFAtMZqL3mpbtl2wfnARHln-dVBag4xZ23Y,4813
 bayinx/dists/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 bayinx/dists/bernoulli.py,sha256=xMV9BgtVX_1XkPdZ43q0meMIEkgMyuUPx--dyo6_DKs,1006
-bayinx/dists/binomial.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-bayinx/dists/gamma.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-bayinx/dists/gamma2.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-bayinx/dists/normal.py,sha256=3CXSgHWnuglmP8cKVUh2Yt4Rb9_LR_mwPRXDm_LuSRo,2679
+bayinx/dists/gamma2.py,sha256=jc9jiNuuIAv3tdWghQ2Y2ANpJbDscPJKOvs6dOJVJD0,1315
+bayinx/dists/normal.py,sha256=AXsf3Xe2BfCyzfcQ5i8J9AD92LedQjgqviUMQd697D8,2628
 bayinx/dists/uniform.py,sha256=mogFe8VuDelM9KXE6RxGek0-tuZYFrwmo_oMOPHXleA,2359
 bayinx/mhx/__init__.py,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
 bayinx/mhx/vi/__init__.py,sha256=YfkXKsqo9Dk_AmQGjZKm4vfG8eLer2ez92G-cOExphs,193
@@ -18,10 +16,10 @@ bayinx/mhx/vi/meanfield.py,sha256=8hM1KZ52TpRPLwiQcowsJLlQ-5nJzUEcKrtDiGrFoSs,37
 bayinx/mhx/vi/normalizing_flow.py,sha256=FvxDtqGRtaEeeF-bXCYnIEAvOOXVHKUK0oCTF9ma02Y,4622
 bayinx/mhx/vi/standard.py,sha256=HaJsIz70Qo1Ql2hMQ-GQhcnfWiOGtyxgkOsm_yQaDKI,1718
 bayinx/mhx/vi/flows/__init__.py,sha256=Hn0Wqvvyv8Vr-mFmimwgNKCByxj-fjrlIvdR7tUSolg,180
-bayinx/mhx/vi/flows/fullaffine.py,sha256=Kvaa8epqaqz9tdMCnf9T_-2P3Bh_TkhA6NrilKHY93A,1886
+bayinx/mhx/vi/flows/fullaffine.py,sha256=s-fxgzv84BEqNMnmLt6vtvwJqBzixCC2OwWXLz3IK-w,1940
 bayinx/mhx/vi/flows/planar.py,sha256=WVj-oxcRctuoRA6KJjU63ek1ZgKNG2vI-TLN0QqjtKA,1916
 bayinx/mhx/vi/flows/radial.py,sha256=Obj3SraliawIHmP14F9wRpWt34y3kscY--Izy24eCvM,2499
 bayinx/mhx/vi/flows/sylvester.py,sha256=ppK0BmS_ThvrCEhJiP_-p-kj67TQHSlU_RUZpDbIhsQ,469
-bayinx-0.2.28.dist-info/METADATA,sha256=xe3Wlo3UlD3VuTc42ChwnPTL6lp3BZmxnuf0gnZxWv0,3058
-bayinx-0.2.28.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-bayinx-0.2.28.dist-info/RECORD,,
+bayinx-0.2.29.dist-info/METADATA,sha256=kfq898UxqnChe-87HPe46_icfS2btMbJPp3INv3IXo0,3058
+bayinx-0.2.29.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+bayinx-0.2.29.dist-info/RECORD,,