efax 2.2.4__tar.gz → 2.3.0__tar.gz

{efax-2.2.4 → efax-2.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: efax
-Version: 2.2.4
+Version: 2.3.0
 Summary: Exponential families for JAX
 Author: Neil Girdhar
 Author-email: Neil Girdhar <mistersheik@gmail.com>

{efax-2.2.4 → efax-2.3.0}/efax/__init__.py

@@ -87,15 +87,26 @@ from ._src.parameter.support import (
     ScalarSupport,
     SimplexSupport,
     SquareMatrixSupport,
+    SubsimplexSupport,
     Support,
     SymmetricMatrixSupport,
     VectorSupport,
 )
 from ._src.parametrization import Distribution, SimpleDistribution
-from ._src.structure.assembler import Assembler, SubDistributionInfo
+from ._src.structure.assembler import (
+    Assembler,
+    JointDistributionInfo,
+    SimpleDistributionInfo,
+    SubDistributionInfo,
+)
 from ._src.structure.estimator import Estimator
 from ._src.structure.flattener import Flattener
-from ._src.tools import parameter_dot_product, parameter_map, parameter_mean
+from ._src.tools import (
+    parameter_dot_product,
+    parameter_holomorphic_dot,
+    parameter_map,
+    parameter_mean,
+)
 from ._src.transform.joint import JointDistribution, JointDistributionE, JointDistributionN
 
 __all__ = [
@@ -152,6 +163,7 @@ __all__ = [
     "IsotropicNormalNP",
     "JointDistribution",
     "JointDistributionE",
+    "JointDistributionInfo",
     "JointDistributionN",
     "LogNormalEP",
     "LogNormalNP",
@@ -184,11 +196,13 @@ __all__ = [
     "Samplable",
     "ScalarSupport",
     "SimpleDistribution",
+    "SimpleDistributionInfo",
     "SimplexSupport",
     "SoftplusNormalEP",
     "SoftplusNormalNP",
     "SquareMatrixSupport",
     "SubDistributionInfo",
+    "SubsimplexSupport",
     "Support",
     "SymmetricMatrixSupport",
     "UnitVarianceLogNormalEP",
@@ -210,6 +224,7 @@ __all__ = [
     "flat_dict_of_parameters",
     "flatten_mapping",
     "parameter_dot_product",
+    "parameter_holomorphic_dot",
     "parameter_map",
     "parameter_mean",
     "parameters",

{efax-2.2.4 → efax-2.3.0}/efax/_src/distributions/categorical.py

@@ -78,9 +78,9 @@ class CategoricalNP(
         if shape is not None:
             shape += self.shape
         logits = xp.concat((self.log_odds, xp.zeros((*self.shape, 1))), axis=-1)
-        retval = xpx.one_hot(jr.categorical(key, logits, shape=shape), self.dimensions() + 1)  # type: ignore
+        retval = xpx.one_hot(jr.categorical(key, logits, shape=shape), self.dimensions() + 1)  # ty: ignore
         assert isinstance(retval, JaxArray)
-        return retval[..., :-1]  # pyright: ignore
+        return retval[..., :-1]
 
     @override
     def dimensions(self) -> int:

{efax-2.2.4 → efax-2.3.0}/efax/_src/distributions/gen_dirichlet.py

@@ -12,14 +12,20 @@ from typing import override
 
 import jax.scipy.special as jss
 from array_api_compat import array_namespace
-from tjax import JaxArray, JaxRealArray, Shape, softplus
+from tjax import JaxArray, JaxRealArray, Shape, inverse_softplus, jit, softplus
 from tjax.dataclasses import dataclass
 
 from efax._src.interfaces.multidimensional import Multidimensional
 from efax._src.mixins.exp_to_nat.exp_to_nat import ExpToNat
 from efax._src.mixins.has_entropy import HasEntropyEP, HasEntropyNP
 from efax._src.natural_parametrization import NaturalParametrization
-from efax._src.parameter import RealField, VectorSupport, distribution_parameter, negative_support
+from efax._src.parameter import (
+    RealField,
+    SubsimplexSupport,
+    VectorSupport,
+    distribution_parameter,
+    negative_support,
+)
 
 
 @dataclass
@@ -42,8 +48,8 @@ class GeneralizedDirichletNP(
 
     @override
     @classmethod
-    def domain_support(cls) -> VectorSupport:
-        return VectorSupport()
+    def domain_support(cls) -> SubsimplexSupport:
+        return SubsimplexSupport()
 
     @override
     def log_normalizer(self) -> JaxRealArray:
@@ -119,14 +125,56 @@ class GeneralizedDirichletEP(
 
     @override
     @classmethod
-    def domain_support(cls) -> VectorSupport:
-        return VectorSupport()
+    def domain_support(cls) -> SubsimplexSupport:
+        return SubsimplexSupport()
 
     @classmethod
     @override
     def natural_parametrization_cls(cls) -> type[GeneralizedDirichletNP]:
         return GeneralizedDirichletNP
 
+    @jit
+    @override
+    def to_nat(self) -> GeneralizedDirichletNP:
+        xp = array_namespace(self)
+        zero = xp.zeros_like(self.mean_log_cumulative_probability[..., :1])
+        beta_bar = xp.concat(
+            [
+                self.mean_log_cumulative_probability[..., :1],
+                xp.diff(self.mean_log_cumulative_probability, axis=-1),
+            ],
+            axis=-1,
+        )
+        alpha_bar_direct = self.mean_log_probability - xp.concat(
+            [zero, self.mean_log_cumulative_probability[..., :-1]],
+            axis=-1,
+        )
+        alpha, beta = _beta_parameters_from_expected_logs(alpha_bar_direct, beta_bar)
+        alpha_roll = xp.concat([alpha[..., 1:], zero], axis=-1)
+        gamma = -xp.diff(beta, axis=-1, append=xp.ones_like(zero)) - alpha_roll
+        return GeneralizedDirichletNP(alpha - 1.0, gamma)
+
+    @override
+    def initial_search_parameters(self) -> JaxRealArray:
+        xp = array_namespace(self)
+        beta_bar = xp.concat(
+            [
+                self.mean_log_cumulative_probability[..., :1],
+                xp.diff(self.mean_log_cumulative_probability, axis=-1),
+            ],
+            axis=-1,
+        )
+        zero = xp.zeros_like(self.mean_log_cumulative_probability[..., :1])
+        alpha_bar_direct = self.mean_log_probability - xp.concat(
+            [zero, self.mean_log_cumulative_probability[..., :-1]],
+            axis=-1,
+        )
+        alpha, beta = _beta_parameters_from_expected_logs(alpha_bar_direct, beta_bar)
+        alpha_roll = xp.concat([alpha[..., 1:], zero], axis=-1)
+        gamma = -xp.diff(beta, axis=-1, append=xp.ones_like(zero)) - alpha_roll
+        gamma = xp.maximum(gamma, 1e-3)
+        return xp.concat([inverse_softplus(alpha), inverse_softplus(gamma)], axis=-1)
+
     @override
     def search_to_natural(self, search_parameters: JaxRealArray) -> GeneralizedDirichletNP:
         # Run Newton's method on the whole real hyperspace.
@@ -144,3 +192,68 @@ class GeneralizedDirichletEP(
     @override
     def dimensions(self) -> int:
         return self.mean_log_probability.shape[-1]
+
+
+inverse_digamma_switch = -2.22
+
+
+def _inverse_digamma(y: JaxRealArray) -> JaxRealArray:
+    xp = array_namespace(y)
+    x = xp.where(
+        y >= inverse_digamma_switch,
+        xp.exp(y) + 0.5,
+        -1.0 / (y - jss.digamma(1.0)),
+    )
+    for _ in range(8):
+        x -= (jss.digamma(x) - y) / jss.polygamma(1, x)
+    return x
+
+
+def _beta_parameters_from_expected_logs(
+    mean_log_probability: JaxRealArray,
+    mean_log_complement: JaxRealArray,
+) -> tuple[JaxRealArray, JaxRealArray]:
+    xp = array_namespace(mean_log_probability, mean_log_complement)
+    mean = xp.exp(mean_log_probability)
+    complement = xp.exp(mean_log_complement)
+    simplex_mean = mean / (mean + complement)
+    concentration = xp.maximum(2.0, 1.0 / xp.maximum(simplex_mean * (1.0 - simplex_mean), 1e-3))
+    for _ in range(20):
+        psi_sum = jss.digamma(concentration)
+        alpha = _inverse_digamma(mean_log_probability + psi_sum)
+        beta = _inverse_digamma(mean_log_complement + psi_sum)
+        concentration = xp.maximum(alpha + beta, 1e-6)
+    alpha = xp.maximum(alpha, 1e-6)
+    beta = xp.maximum(beta, 1e-6)
+    log_alpha = xp.log(alpha)
+    log_beta = xp.log(beta)
+    for _ in range(20):
+        alpha = xp.exp(log_alpha)
+        beta = xp.exp(log_beta)
+        delta_log_alpha, delta_log_beta = _beta_expected_log_newton_step(
+            alpha, beta, mean_log_probability, mean_log_complement
+        )
+        log_alpha -= delta_log_alpha
+        log_beta -= delta_log_beta
+    alpha = xp.exp(log_alpha)
+    beta = xp.exp(log_beta)
+    return alpha, beta
+
+
+def _beta_expected_log_newton_step(
+    alpha: JaxRealArray,
+    beta: JaxRealArray,
+    mean_log_probability: JaxRealArray,
+    mean_log_complement: JaxRealArray,
+) -> tuple[JaxRealArray, JaxRealArray]:
+    trigamma_sum = jss.polygamma(1, alpha + beta)
+    f1 = jss.digamma(alpha) - jss.digamma(alpha + beta) - mean_log_probability
+    f2 = jss.digamma(beta) - jss.digamma(alpha + beta) - mean_log_complement
+    j11 = alpha * (jss.polygamma(1, alpha) - trigamma_sum)
+    j12 = -beta * trigamma_sum
+    j21 = -alpha * trigamma_sum
+    j22 = beta * (jss.polygamma(1, beta) - trigamma_sum)
+    determinant = j11 * j22 - j12 * j21
+    delta_log_alpha = (j22 * f1 - j12 * f2) / determinant
+    delta_log_beta = (-j21 * f1 + j11 * f2) / determinant
+    return delta_log_alpha, delta_log_beta

{efax-2.2.4 → efax-2.3.0}/efax/_src/distributions/log_normal/log_normal.py

@@ -42,7 +42,7 @@ class LogNormalNP(
     @override
     @classmethod
     def domain_support(cls) -> ScalarSupport:
-        return ScalarSupport()
+        return ScalarSupport(ring=positive_support)
 
     @override
     @classmethod
@@ -103,7 +103,7 @@ class LogNormalEP(
     @override
     @classmethod
     def domain_support(cls) -> ScalarSupport:
-        return ScalarSupport()
+        return ScalarSupport(ring=positive_support)
 
     @classmethod
     @override

{efax-2.2.4 → efax-2.3.0}/efax/_src/distributions/log_normal/unit_variance.py

@@ -17,7 +17,7 @@ from efax._src.mixins.transformed_parametrization import (
     TransformedNaturalParametrization,
 )
 from efax._src.natural_parametrization import NaturalParametrization
-from efax._src.parameter import ScalarSupport, distribution_parameter
+from efax._src.parameter import ScalarSupport, distribution_parameter, positive_support
 
 
 @dataclass
@@ -35,7 +35,7 @@ class UnitVarianceLogNormalNP(
     @override
     @classmethod
     def domain_support(cls) -> ScalarSupport:
-        return ScalarSupport()
+        return ScalarSupport(ring=positive_support)
 
     @override
     @classmethod
@@ -90,7 +90,7 @@ class UnitVarianceLogNormalEP(
     @override
     @classmethod
     def domain_support(cls) -> ScalarSupport:
-        return ScalarSupport()
+        return ScalarSupport(ring=positive_support)
 
     @classmethod
     @override

{efax-2.2.4 → efax-2.3.0}/efax/_src/distributions/multivariate_normal/fixed_variance.py

@@ -76,7 +76,10 @@ class MultivariateFixedVarianceNormalNP(
     def sufficient_statistics(
         cls, x: JaxRealArray, **fixed_parameters: JaxArray
     ) -> MultivariateFixedVarianceNormalEP:
-        return MultivariateFixedVarianceNormalEP(x, variance=fixed_parameters["variance"])
+        xp = array_namespace(x)
+        return MultivariateFixedVarianceNormalEP(
+            x, variance=xp.broadcast_to(fixed_parameters["variance"], x.shape[:-1])
+        )
 
     @override
     def sample(self, key: KeyArray, shape: Shape | None = None) -> JaxRealArray:

{efax-2.2.4 → efax-2.3.0}/efax/_src/distributions/weibull.py

@@ -68,8 +68,8 @@ class WeibullNP(
     @classmethod
     def sufficient_statistics(cls, x: JaxRealArray, **fixed_parameters: JaxArray) -> WeibullEP:
         xp = array_namespace(x)
-        concentration = fixed_parameters["concentration"]
-        return WeibullEP(xp.broadcast_to(concentration, x.shape), x**concentration)
+        concentration = xp.broadcast_to(fixed_parameters["concentration"], x.shape)
+        return WeibullEP(concentration, x**concentration)
 
     @override
     def sample(self, key: KeyArray, shape: Shape | None = None) -> JaxRealArray:

{efax-2.2.4 → efax-2.3.0}/efax/_src/mixins/exp_to_nat/exp_to_nat.py

@@ -43,7 +43,7 @@ class ExpToNat(ExpectationParametrization[NP], SimpleDistribution, Generic[NP]):
         # Select a default minimizer based on the dimensionality of the search space.
         # Bisection is used for scalar problems; Newton's method for vector problems.
         if hasattr(super(), "__post_init__"):
-            super().__post_init__()  # pyright: ignore
+            super().__post_init__()  # ty: ignore
         if self.minimizer is None:
             from .optimistix import default_bisection_minimizer, default_minimizer  # noqa: PLC0415
 
@@ -96,7 +96,7 @@ class ExpToNat(ExpectationParametrization[NP], SimpleDistribution, Generic[NP]):
             unflatten_as_type=np_cls,
             mapped_to_plane=True,
         )
-        return flattener.unflatten(search_parameters)  # type: ignore
+        return flattener.unflatten(search_parameters)  # ty: ignore
 
     def search_gradient(self, search_parameters: SP) -> SP:
         """Convert the search parameters to the natural gradient.

{efax-2.2.4 → efax-2.3.0}/efax/_src/mixins/exp_to_nat/optimistix.py

@@ -55,13 +55,13 @@ class OptimistixRootFinder(ExpToNatMinimizer):
 
 
 default_minimizer = OptimistixRootFinder(
-    solver=optx.Newton[JaxRealArray, JaxRealArray, None](rtol=0.0, atol=1e-7),  # ty: ignore
+    solver=optx.Newton[JaxRealArray, JaxRealArray, None](rtol=0.0, atol=1e-7),
     max_steps=1000,
 )
 
 
 default_bisection_minimizer = OptimistixRootFinder(
-    solver=optx.Bisection(  # type: ignore
+    solver=optx.Bisection(  # ty: ignore
         rtol=0.0, atol=1e-7, flip="detect", expand_if_necessary=True
     ),
     max_steps=1000,

{efax-2.2.4 → efax-2.3.0}/efax/_src/natural_parametrization.py

@@ -24,7 +24,7 @@ from .parametrization import Distribution
 from .structure.assembler import Assembler
 from .structure.estimator import Estimator
 from .structure.flattener import Flattener
-from .tools import parameter_dot_product
+from .tools import parameter_dot_product, parameter_holomorphic_dot
 
 if TYPE_CHECKING:
     from .expectation_parametrization import ExpectationParametrization
@@ -37,13 +37,24 @@ Domain = TypeVar("Domain", bound=JaxComplexArray | dict[str, Any], default=Any)
 def _log_normalizer_jvp(
     primals: tuple[NaturalParametrization],
     tangents: tuple[NaturalParametrization],
-) -> tuple[JaxRealArray, JaxRealArray]:
-    """The log-normalizer's special JVP vastly improves numerical stability."""
+) -> tuple[JaxArray, JaxArray]:
+    """The log-normalizer's special JVP vastly improves numerical stability.
+
+    On the analytic-continuation path used by ``characteristic_function`` the
+    natural parameters live in ℂⁿ and the primal output ``y`` is complex; the
+    holomorphic derivative ``Σ q_dot · p`` is what JAX expects as the complex
+    tangent.  In every other case (real distributions, and intrinsically-complex
+    distributions whose ``log_normalizer`` returns a real scalar) the Hermitian
+    form is correct.  Dispatching on ``y.dtype`` distinguishes the two.
+    """
     (q,) = primals
     (q_dot,) = tangents
     y = q.log_normalizer()
     p = q.to_exp()
-    y_dot = parameter_dot_product(q_dot, p)
+    if jnp.iscomplexobj(y):
+        y_dot = parameter_holomorphic_dot(q_dot, p)
+    else:
+        y_dot = parameter_dot_product(q_dot, p)
     return y, y_dot
 
 
@@ -157,7 +168,7 @@ class NaturalParametrization(Distribution, JaxAbstractClass, Generic[EP, Domain]
         """
         xp = array_namespace(self)
         fisher_information_diagonal = self.fisher_information_diagonal()
-        structure = Assembler.create_assembler(self)
+        assembler = Assembler.create_assembler(self)
         final_parameters = parameters(self)
         for path, (value, support) in parameters(fisher_information_diagonal, support=True).items():
             na = support.axes()
@@ -170,7 +181,7 @@ class NaturalParametrization(Distribution, JaxAbstractClass, Generic[EP, Domain]
             else:
                 raise RuntimeError
             final_parameters[path] = new_value
-        return structure.assemble(final_parameters)
+        return assembler.assemble(final_parameters)
 
     @final
     def jeffreys_prior_density(self) -> JaxRealArray:

{efax-2.2.4 → efax-2.3.0}/efax/_src/parameter/__init__.py

@@ -17,6 +17,7 @@ from .support import (
     ScalarSupport,
     SimplexSupport,
     SquareMatrixSupport,
+    SubsimplexSupport,
     Support,
     SymmetricMatrixSupport,
     VectorSupport,
@@ -32,6 +33,7 @@ __all__ = [
     "ScalarSupport",
     "SimplexSupport",
     "SquareMatrixSupport",
+    "SubsimplexSupport",
     "Support",
     "SymmetricMatrixSupport",
     "VectorSupport",

{efax-2.2.4 → efax-2.3.0}/efax/_src/parameter/ring.py

@@ -16,6 +16,7 @@ from tjax import (
     JaxRealArray,
     RealNumeric,
     Shape,
+    divide_where,
     inverse_softplus,
     softplus,
 )
@@ -142,7 +143,7 @@ class RealField(Ring):
             minimum = xp.asarray(self.minimum)
             maximum = xp.asarray(self.maximum)
             domain_width = xp.maximum(maximum - minimum - safety, 0.0)
-            true_minimum = xp.mean(minimum, maximum) - domain_width * 0.5
+            true_minimum = (minimum + maximum - domain_width) * 0.5
             return xp.asarray(true_minimum + domain_width * rng.uniform(size=shape))
         return xp.asarray(rng.normal(scale=1.0, size=shape) * self.generation_scale)
 
@@ -176,12 +177,22 @@ class ComplexField(Ring):
                 # x is outside the disk of the given minimum.  Map it to the plane.
                 magnitude = xp.abs(x)
                 corrected_magnitude = magnitude - minimum
-                corrected_x = x * corrected_magnitude / magnitude
+                corrected_x = x * divide_where(
+                    corrected_magnitude,
+                    magnitude,
+                    where=magnitude != 0.0,
+                    otherwise=xp.zeros_like(magnitude),
+                )
             else:
                 # x is in the disk of the given maximum.  Map it to the plane.
                 magnitude = xp.abs(x)
                 corrected_magnitude = jss.logit(((magnitude - minimum) / maximum) * 0.5 + 0.5)
-                corrected_x = x * corrected_magnitude / magnitude
+                corrected_x = x * divide_where(
+                    corrected_magnitude,
+                    magnitude,
+                    where=magnitude != 0.0,
+                    otherwise=xp.zeros_like(magnitude),
+                )
         else:
             corrected_x = x
         return xp.concat([xp.real(corrected_x), xp.imag(corrected_x)], axis=-1)
@@ -206,11 +217,21 @@ class ComplexField(Ring):
             # x is outside the disk of the given minimum.  Map it to the plane.
             corrected_magnitude = xp.abs(corrected_x)
             magnitude = corrected_magnitude + minimum
-            return corrected_x * magnitude / corrected_magnitude
+            return corrected_x * divide_where(
+                magnitude,
+                corrected_magnitude,
+                where=corrected_magnitude != 0.0,
+                otherwise=xp.zeros_like(corrected_magnitude),
+            )
         # x is in the disk of the given maximum.  Map it to the plane.
         corrected_magnitude = xp.abs(corrected_x)
         magnitude = maximum * (jss.expit(corrected_magnitude) - 0.5) * 2.0 + minimum
-        return corrected_x * magnitude / corrected_magnitude
+        return corrected_x * divide_where(
+            magnitude,
+            corrected_magnitude,
+            where=corrected_magnitude != 0.0,
+            otherwise=xp.zeros_like(corrected_magnitude),
+        )
 
     @override
     def generate(self, xp: Namespace, rng: Generator, shape: Shape, safety: float) -> JaxRealArray:

{efax-2.2.4 → efax-2.3.0}/efax/_src/parameter/support.py

@@ -10,7 +10,7 @@ import numpy as np
 from array_api_compat import array_namespace
 from numpy.random import Generator
 from opt_einsum import contract
-from tjax import JaxArray, JaxRealArray, Shape
+from tjax import JaxArray, JaxRealArray, Shape, divide_where
 
 from efax._src.types import Namespace
 
@@ -141,13 +141,22 @@ class SimplexSupport(Support):
 
     @override
     def flattened(self, x: JaxArray, *, map_to_plane: bool) -> JaxRealArray:
-        return self.ring.flattened(x, map_to_plane=map_to_plane)
+        if not map_to_plane:
+            return self.ring.flattened(x, map_to_plane=False)
+        xp = array_namespace(x)
+        residual = 1.0 - xp.sum(x, axis=-1, keepdims=True)
+        return xp.log(x / residual)
 
     @override
     def unflattened(self, y: JaxRealArray, dimensions: int, *, map_from_plane: bool) -> JaxArray:
-        x = self.ring.unflattened(y, map_from_plane=map_from_plane)
-        assert x.shape[-1] == dimensions
-        return x
+        if not map_from_plane:
+            x = self.ring.unflattened(y, map_from_plane=False)
+            assert x.shape[-1] == dimensions - 1
+            return x
+        xp = array_namespace(y)
+        assert y.shape[-1] == dimensions - 1
+        logits = xp.concat((y, xp.zeros((*y.shape[:-1], 1), dtype=y.dtype)), axis=-1)
+        return jss.softmax(logits, axis=-1)[..., :-1]
 
     @override
     def clamp(self, x: JaxArray) -> JaxArray:
@@ -161,7 +170,62 @@ class SimplexSupport(Support):
     def generate(
         self, xp: Namespace, rng: Generator, shape: Shape, safety: float, dimensions: int
     ) -> JaxRealArray:
-        raise NotImplementedError
+        alpha = xp.ones(dimensions)
+        x = xp.asarray(rng.dirichlet(alpha, size=shape))
+        if safety > 0.0:
+            x = x * (1.0 - safety) + safety / dimensions
+        return x[..., :-1]
+
+
+class SubsimplexSupport(Support):
+    @override
+    def axes(self) -> int:
+        return 1
+
+    @override
+    def shape(self, dimensions: int) -> Shape:
+        return (dimensions,)
+
+    @override
+    def num_elements(self, dimensions: int) -> int:
+        return self.ring.num_elements(dimensions)
+
+    @override
+    def flattened(self, x: JaxArray, *, map_to_plane: bool) -> JaxRealArray:
+        if not map_to_plane:
+            return self.ring.flattened(x, map_to_plane=False)
+        xp = array_namespace(x)
+        residual = 1.0 - xp.sum(x, axis=-1, keepdims=True)
+        return xp.log(x / residual)
+
+    @override
+    def unflattened(self, y: JaxRealArray, dimensions: int, *, map_from_plane: bool) -> JaxArray:
+        if not map_from_plane:
+            x = self.ring.unflattened(y, map_from_plane=False)
+            assert x.shape[-1] == dimensions
+            return x
+        xp = array_namespace(y)
+        assert y.shape[-1] == dimensions
+        logits = xp.concat((y, xp.zeros((*y.shape[:-1], 1), dtype=y.dtype)), axis=-1)
+        return jss.softmax(logits, axis=-1)[..., :-1]
+
+    @override
+    def clamp(self, x: JaxArray) -> JaxArray:
+        xp = array_namespace(x)
+        eps = xp.finfo(x.dtype).eps
+        s = xp.sum(x, axis=-1, keepdims=True)
+        x *= xp.minimum(1.0, (1.0 - eps) / s)
+        return xp.clip(x, min=eps, max=1.0 - eps)
+
+    @override
+    def generate(
+        self, xp: Namespace, rng: Generator, shape: Shape, safety: float, dimensions: int
+    ) -> JaxRealArray:
+        alpha = xp.ones(dimensions + 1)
+        x = xp.asarray(rng.dirichlet(alpha, size=shape))[..., :-1]
+        if safety > 0.0:
+            x *= 1.0 - safety
+        return x
 
 
 class SymmetricMatrixSupport(Support):
@@ -224,10 +288,10 @@ class SymmetricMatrixSupport(Support):
             i = int(i_)
             j = int(j_)
             xk = x[..., k]
-            result = xpx.at(result)[..., i, j].set(xk)  # type: ignore
+            result = xpx.at(result)[..., i, j].set(xk)  # ty: ignore
             if i != j:
                 cxk = xp.conj(xk) if self.hermitian else xk
-                result = xpx.at(result)[..., j, i].set(cxk)  # type: ignore
+                result = xpx.at(result)[..., j, i].set(cxk)  # ty: ignore
         assert isinstance(result, JaxArray)
         return result
 
@@ -299,7 +363,12 @@ class CircularBoundedSupport(VectorSupport):
         xp = array_namespace(x)
         magnitude = xp.linalg.norm(x, 2, axis=-1, keepdims=True)
         corrected_magnitude = jss.logit((magnitude / self.radius) * 0.5 + 0.5)
-        return x * corrected_magnitude / magnitude
+        return x * divide_where(
+            corrected_magnitude,
+            magnitude,
+            where=magnitude != 0.0,
+            otherwise=xp.zeros_like(magnitude),
+        )
 
     @override
     def unflattened(self, y: JaxRealArray, dimensions: int, *, map_from_plane: bool) -> JaxArray:
@@ -310,4 +379,9 @@ class CircularBoundedSupport(VectorSupport):
         assert y.shape[-1] == dimensions
         corrected_magnitude = cast("JaxRealArray", xp.linalg.norm(y, 2, axis=-1, keepdims=True))
         magnitude = self.radius * (jss.expit(corrected_magnitude) - 0.5) * 2.0
-        return y * magnitude / corrected_magnitude
+        return y * divide_where(
+            magnitude,
+            corrected_magnitude,
+            where=corrected_magnitude != 0.0,
+            otherwise=xp.zeros_like(corrected_magnitude),
+        )