Merge pull request #121 from danielward27/laplace

add Laplace

flowjax/distributions.py

@@ -272,7 +272,7 @@ def _sample(self, key, condition=None):
     def _sample_and_log_prob(
         self,
         key: Array,
-        condition=None,
+        condition: Array | None = None,
     ):  # TODO add overide decorator when python>=3.12 is common
         # We override to avoid computing the inverse transformation.
         base_sample, log_prob_base = self.base_dist._sample_and_log_prob(key, condition)
@@ -408,7 +408,7 @@ class LogNormal(AbstractTransformed):
 
     Args:
         loc: Location paramter. Defaults to 0.
-        scale: Scale parameter. Defaults to 1.0.
+        scale: Scale parameter. Defaults to 1.
     """
 
     base_dist: StandardNormal
@@ -527,7 +527,7 @@ class Gumbel(AbstractTransformed):
 
     Args:
         loc: Location paramter.
-        scale: Scale parameter. Defaults to 1.0.
+        scale: Scale parameter. Defaults to 1.
     """
 
     base_dist: _StandardGumbel
@@ -573,7 +573,7 @@ class Cauchy(AbstractTransformed):
 
     Args:
         loc: Location paramter.
-        scale: Scale parameter. Defaults to 1.0.
+        scale: Scale parameter. Defaults to 1.
     """
 
     base_dist: _StandardCauchy
@@ -628,8 +628,8 @@ class StudentT(AbstractTransformed):
 
     Args:
         df: The degrees of freedom.
-        loc: Location parameter. Defaults to 0.0.
-        scale: Scale parameter. Defaults to 1.0.
+        loc: Location parameter. Defaults to 0.
+        scale: Scale parameter. Defaults to 1.
     """
 
     base_dist: _StandardStudentT
@@ -656,6 +656,48 @@ def df(self):
         return self.base_dist.df
 
 
+class _StandardLaplace(AbstractDistribution):
+    """Implements standard laplace distribution (loc=0, scale=1)."""
+
+    shape: tuple[int, ...] = ()
+    cond_shape: ClassVar[None] = None
+
+    def _log_prob(self, x, condition=None):
+        return jstats.laplace.logpdf(x).sum()
+
+    def _sample(self, key, condition=None):
+        return jr.laplace(key, shape=self.shape)
+
+
+class Laplace(AbstractTransformed):
+    """Laplace distribution.
+
+    ``loc`` and ``scale`` should broadcast to the dimension of the distribution..
+
+    Args:
+        loc: Location paramter. Defaults to 0.
+        scale: Scale parameter. Defaults to 1.
+    """
+
+    base_dist: _StandardLaplace
+    bijection: Affine
+
+    def __init__(self, loc: ArrayLike = 0, scale: ArrayLike = 1):
+        shape = jnp.broadcast_shapes(jnp.shape(loc), jnp.shape(scale))
+        self.base_dist = _StandardLaplace(shape)
+        self.bijection = Affine(loc, scale)
+
+    @property
+    def loc(self):
+        """Location of the distribution."""
+        return self.bijection.loc
+
+    @property
+    def scale(self):
+        """Scale of the distribution."""
+        return self.bijection.scale
+
+
 class SpecializeCondition(AbstractDistribution):  # TODO check tested
     """Specialise a distribution to a particular conditioning variable instance.
 

pyproject.toml

@@ -16,7 +16,7 @@ license = { file = "LICENSE" }
 name = "flowjax"
 readme = "README.md"
 requires-python = ">=3.10"
-version = "11.0.0"
+version = "11.1.0"
 
 [project.urls]
 repository = "https://github.com/danielward27/flowjax"

tests/test_distributions.py

@@ -13,6 +13,7 @@
     AbstractTransformed,
     Cauchy,
     Gumbel,
+    Laplace,
     LogNormal,
     MultivariateNormal,
     Normal,
@@ -22,6 +23,7 @@
     Uniform,
     _StandardCauchy,
     _StandardGumbel,
+    _StandardLaplace,
     _StandardStudentT,
     _StandardUniform,
 )
@@ -44,6 +46,8 @@
     "_StandardStudentT": lambda shape: _StandardStudentT(jnp.ones(shape)),
     "StudentT": lambda shape: StudentT(jnp.ones(shape)),
     "LogNormal": lambda shape: LogNormal(jnp.ones(shape), 2),
+    "_StandardLaplace": _StandardLaplace,
+    "Laplace": lambda shape: Laplace(jnp.ones(shape)),
 }