efax 2.2.0__tar.gz → 2.2.4__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: efax
-Version: 2.2.0
+Version: 2.2.4
 Summary: Exponential families for JAX
 Author: Neil Girdhar
 Author-email: Neil Girdhar <mistersheik@gmail.com>
@@ -23,7 +23,7 @@ Requires-Dist: numpy>=2.0
 Requires-Dist: opt-einsum>=3.4
 Requires-Dist: optimistix>=0.0.9
 Requires-Dist: optype[numpy]>=0.8.0
-Requires-Dist: tjax>=1.4.4
+Requires-Dist: tjax>=1.4.6
 Requires-Dist: typing-extensions>=4.8
 Maintainer: Neil Girdhar
 Maintainer-email: Neil Girdhar <mistersheik@gmail.com>
@@ -50,6 +50,12 @@ Description-Content-Type: text/x-rst
    :target: https://scientific-python.org/specs/spec-0000/
    :alt: SPEC-0
    :align: center
+.. image:: https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json
+   :alt: Ruff
+   :target: https://github.com/astral-sh/ruff
+.. image:: https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ty/main/assets/badge/v0.json
+   :alt: ty
+   :target: https://github.com/astral-sh/ty
 .. image:: https://img.shields.io/pypi/pyversions/efax
    :alt: PyPI - Python Version
    :align: center
@@ -161,7 +167,10 @@ Every `NaturalParametrization` has methods:
 - `jeffreys_prior_density`, which returns the square root of the Fisher information
   determinant,
 - `characteristic_function`, which evaluates the characteristic function of the sufficient
-  statistics via analytic continuation of the log-normalizer, and
+  statistics via analytic continuation of the log-normalizer (note: for distributions such
+  as `GammaNP` that keep one or more natural parameters real in `_complexify`, the result
+  is approximate — check `NaturalParametrization.characteristic_function_exact` on the
+  class before relying on every component), and
 - `kl_divergence`, which is the KL divergence.
 
 Every `ExpectationParametrization` has methods:
@@ -169,6 +178,11 @@ Every `ExpectationParametrization` has methods:
 - `to_nat` to convert itself to natural parameters, and
 - `kl_divergence`, which is the KL divergence.
 
+The module-level function `expectation_parameters_from_characteristic_function(t, cf_values)`
+estimates the expectation parameters from characteristic-function evaluations via ordinary
+least squares.  It is exact for the normal distribution (whose log-CF is linear in the
+frequencies) and a first-order approximation for other families.
+
 Some parametrizations inherit from these interfaces:
 
 - `HasConjugatePrior` can produce and recover the conjugate prior,
@@ -313,7 +327,7 @@ EFAX supports the following distributions:
   - Dirichlet
   - generalized Dirichlet
 
-- on the n-sphere:
+- on spheres:
 
   - von Mises-Fisher
   - complex von Mises
@@ -448,17 +462,23 @@ Using the cross entropy to iteratively optimize a prediction is simple:
         return x - 1e-4 * x_bar
 
 
-    def body_fun(q: BernoulliNP) -> BernoulliNP:
-        q_bar = gradient_cross_entropy(target_distribution, q)
-        return parameter_map(apply, q, q_bar)
+    # Carry (q, q_bar) as the loop state so the gradient is computed exactly
+    # once per step rather than once in cond_fun and once in body_fun.
+    type State = tuple[BernoulliNP, BernoulliNP]
 
 
-    def cond_fun(q: BernoulliNP) -> JaxBooleanArray:
-        q_bar = gradient_cross_entropy(target_distribution, q)
+    def cond_fun(state: State) -> JaxBooleanArray:
+        _, q_bar = state
         total = jnp.sum(parameter_dot_product(q_bar, q_bar))
         return total > 1e-6  # noqa: PLR2004
 
 
+    def body_fun(state: State) -> State:
+        q, q_bar = state
+        q_new = parameter_map(apply, q, q_bar)
+        return q_new, gradient_cross_entropy(target_distribution, q_new)
+
+
     # The target_distribution is represented as the expectation parameters of a
     # Bernoulli distribution corresponding to probabilities 0.3, 0.4, and 0.7.
     target_distribution = BernoulliEP(jnp.asarray([0.3, 0.4, 0.7]))
@@ -467,10 +487,13 @@ Using the cross entropy to iteratively optimize a prediction is simple:
     # of a Bernoulli distribution corresponding to log-odds 0, which is probability
     # 0.5.
     initial_predictive_distribution = BernoulliNP(jnp.zeros(3))
+    initial_gradient = gradient_cross_entropy(target_distribution,
+                                              initial_predictive_distribution)
 
     # Optimize the predictive distribution iteratively.
-    predictive_distribution = lax.while_loop(cond_fun, body_fun,
-                                             initial_predictive_distribution)
+    predictive_distribution, _ = lax.while_loop(
+        cond_fun, body_fun, (initial_predictive_distribution, initial_gradient)
+    )
 
     # Compare the optimized predictive distribution with the target value in the
     # same natural parametrization.
@@ -535,6 +558,8 @@ instead.
     # Dirichlet distribution.
 
     # Take the mean over the first axis.
+    # parameter_mean averages only variable parameters; fixed parameters (e.g.
+    # the failure count in NegativeBinomialNP) are preserved unchanged.
     ss_mean = parameter_mean(ss, axis=0)  # ss_mean also has type DirichletEP.
 
     # Convert this back to the natural parametrization.

{efax-2.2.0 → efax-2.2.4}/README.rst

@@ -17,6 +17,12 @@
    :target: https://scientific-python.org/specs/spec-0000/
    :alt: SPEC-0
    :align: center
+.. image:: https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json
+   :alt: Ruff
+   :target: https://github.com/astral-sh/ruff
+.. image:: https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ty/main/assets/badge/v0.json
+   :alt: ty
+   :target: https://github.com/astral-sh/ty
 .. image:: https://img.shields.io/pypi/pyversions/efax
    :alt: PyPI - Python Version
    :align: center
@@ -128,7 +134,10 @@ Every `NaturalParametrization` has methods:
 - `jeffreys_prior_density`, which returns the square root of the Fisher information
   determinant,
 - `characteristic_function`, which evaluates the characteristic function of the sufficient
-  statistics via analytic continuation of the log-normalizer, and
+  statistics via analytic continuation of the log-normalizer (note: for distributions such
+  as `GammaNP` that keep one or more natural parameters real in `_complexify`, the result
+  is approximate — check `NaturalParametrization.characteristic_function_exact` on the
+  class before relying on every component), and
 - `kl_divergence`, which is the KL divergence.
 
 Every `ExpectationParametrization` has methods:
@@ -136,6 +145,11 @@ Every `ExpectationParametrization` has methods:
 - `to_nat` to convert itself to natural parameters, and
 - `kl_divergence`, which is the KL divergence.
 
+The module-level function `expectation_parameters_from_characteristic_function(t, cf_values)`
+estimates the expectation parameters from characteristic-function evaluations via ordinary
+least squares.  It is exact for the normal distribution (whose log-CF is linear in the
+frequencies) and a first-order approximation for other families.
+
 Some parametrizations inherit from these interfaces:
 
 - `HasConjugatePrior` can produce and recover the conjugate prior,
@@ -280,7 +294,7 @@ EFAX supports the following distributions:
   - Dirichlet
   - generalized Dirichlet
 
-- on the n-sphere:
+- on spheres:
 
   - von Mises-Fisher
   - complex von Mises
@@ -415,17 +429,23 @@ Using the cross entropy to iteratively optimize a prediction is simple:
         return x - 1e-4 * x_bar
 
 
-    def body_fun(q: BernoulliNP) -> BernoulliNP:
-        q_bar = gradient_cross_entropy(target_distribution, q)
-        return parameter_map(apply, q, q_bar)
+    # Carry (q, q_bar) as the loop state so the gradient is computed exactly
+    # once per step rather than once in cond_fun and once in body_fun.
+    type State = tuple[BernoulliNP, BernoulliNP]
 
 
-    def cond_fun(q: BernoulliNP) -> JaxBooleanArray:
-        q_bar = gradient_cross_entropy(target_distribution, q)
+    def cond_fun(state: State) -> JaxBooleanArray:
+        _, q_bar = state
         total = jnp.sum(parameter_dot_product(q_bar, q_bar))
         return total > 1e-6  # noqa: PLR2004
 
 
+    def body_fun(state: State) -> State:
+        q, q_bar = state
+        q_new = parameter_map(apply, q, q_bar)
+        return q_new, gradient_cross_entropy(target_distribution, q_new)
+
+
     # The target_distribution is represented as the expectation parameters of a
     # Bernoulli distribution corresponding to probabilities 0.3, 0.4, and 0.7.
     target_distribution = BernoulliEP(jnp.asarray([0.3, 0.4, 0.7]))
@@ -434,10 +454,13 @@ Using the cross entropy to iteratively optimize a prediction is simple:
     # of a Bernoulli distribution corresponding to log-odds 0, which is probability
     # 0.5.
     initial_predictive_distribution = BernoulliNP(jnp.zeros(3))
+    initial_gradient = gradient_cross_entropy(target_distribution,
+                                              initial_predictive_distribution)
 
     # Optimize the predictive distribution iteratively.
-    predictive_distribution = lax.while_loop(cond_fun, body_fun,
-                                             initial_predictive_distribution)
+    predictive_distribution, _ = lax.while_loop(
+        cond_fun, body_fun, (initial_predictive_distribution, initial_gradient)
+    )
 
     # Compare the optimized predictive distribution with the target value in the
     # same natural parametrization.
@@ -502,6 +525,8 @@ instead.
     # Dirichlet distribution.
 
     # Take the mean over the first axis.
+    # parameter_mean averages only variable parameters; fixed parameters (e.g.
+    # the failure count in NegativeBinomialNP) are preserved unchanged.
     ss_mean = parameter_mean(ss, axis=0)  # ss_mean also has type DirichletEP.
 
     # Convert this back to the natural parametrization.

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

@@ -108,7 +108,8 @@ class CategoricalEP(
     """The expectation parametrization of the categorical distribution.
 
     Args:
-        probability: The probability vector with the final element omitted, i.e., [p_i]_{i in 1...n-1}.
+        probability: The probability vector with the final element omitted, i.e.,
+            [p_i]_{i in 1...n-1}.
     """
 
     probability: JaxRealArray = distribution_parameter(

{efax-2.2.0 → efax-2.2.4}/efax/_src/distributions/gamma.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Self, override
+from typing import ClassVar, Self, override
 
 import jax.random as jr
 import jax.scipy.special as jss
@@ -39,6 +39,8 @@ class GammaNP(
         shape_minus_one: The shape minus one.
     """
 
+    characteristic_function_exact: ClassVar[bool] = False
+
     negative_rate: JaxRealArray = distribution_parameter(ScalarSupport(ring=negative_support))
     shape_minus_one: JaxRealArray = distribution_parameter(
         ScalarSupport(ring=RealField(minimum=-1.0))
@@ -83,9 +85,10 @@ class GammaNP(
     def _complexify(self, t: Self) -> Self:  # ty: ignore
         """Only complexify negative_rate; jss.gammaln does not accept complex inputs.
 
-        This means characteristic_function cannot query the log x sufficient
-        statistic (the shape_minus_one component).  Querying it returns 1
-        instead of the true E[exp(i·t·log x)].
+        shape_minus_one is kept real, so ``characteristic_function`` cannot
+        query the log-x sufficient statistic: that component returns 1 instead
+        of the true E[exp(i·t·log x)].  Consequently,
+        ``characteristic_function_exact = False`` on this class.
         """
         return type(self)(self.negative_rate + 1j * t.negative_rate, self.shape_minus_one)
 

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

@@ -55,7 +55,8 @@ class OptimistixRootFinder(ExpToNatMinimizer):
 
 
 default_minimizer = OptimistixRootFinder(
-    solver=optx.Newton[JaxRealArray, JaxRealArray, None](rtol=0.0, atol=1e-7), max_steps=1000
+    solver=optx.Newton[JaxRealArray, JaxRealArray, None](rtol=0.0, atol=1e-7),  # ty: ignore
+    max_steps=1000,
 )
 
 

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

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from abc import abstractmethod
-from typing import TYPE_CHECKING, Any, Generic, Self, final, get_type_hints
+from typing import TYPE_CHECKING, Any, ClassVar, Generic, Self, final, get_type_hints
 
 import jax
 import jax.numpy as jnp
@@ -53,8 +53,18 @@ class NaturalParametrization(Distribution, JaxAbstractClass, Generic[EP, Domain]
 
     The motivation for the natural parametrization is combining and scaling independent predictive
     evidence.  In the natural parametrization, these operations correspond to scaling and addition.
+
+    Class variables:
+        characteristic_function_exact: True when `_complexify` shifts every natural parameter
+            into the complex plane, so `characteristic_function` queries all sufficient
+            statistics exactly.  Set to False in subclasses that must keep one or more
+            parameters real (e.g. because `log_normalizer` calls a function such as
+            ``gammaln`` that does not accept complex inputs); those components return 1
+            instead of the true characteristic-function value.
     """
 
+    characteristic_function_exact: ClassVar[bool] = True
+
     @abstract_custom_jvp(_log_normalizer_jvp)
     @abstract_jit
     @abstractmethod
@@ -196,11 +206,13 @@ class NaturalParametrization(Distribution, JaxAbstractClass, Generic[EP, Domain]
     def _complexify(self, t: Self) -> Self:
         """Shift natural parameters into the complex plane: η → η + i·t.
 
-        The default shifts every field.  Subclasses must override this when
-        ``log_normalizer`` calls functions that don't accept complex inputs
-        (e.g. ``gammaln``).  Keeping parameter k real means
-        ``characteristic_function`` cannot query the k-th sufficient statistic:
-        it will silently return 1 for that component instead of the true value.
+        The default shifts every field, which is correct for any distribution
+        whose ``log_normalizer`` accepts fully complex inputs.  Subclasses that
+        must keep one or more parameters real (e.g. because ``log_normalizer``
+        calls ``gammaln``, which rejects complex inputs) must override this
+        method **and** set ``characteristic_function_exact = False`` on the
+        class.  For those parameters, ``characteristic_function`` returns 1
+        instead of the true E[exp(i·⟨t, T(x)⟩)] component.
 
         Args:
             t: Imaginary displacements, same pytree structure as self.
@@ -237,6 +249,13 @@ class NaturalParametrization(Distribution, JaxAbstractClass, Generic[EP, Domain]
             t_grid = PoissonNP(omegas)           # shape (k,)
             phi = p.characteristic_function(t_grid)  # shape (k,)
 
+        Note: when ``type(self).characteristic_function_exact`` is False, one
+        or more natural parameters are kept real in ``_complexify`` because
+        ``log_normalizer`` uses a function that rejects complex inputs.  The
+        corresponding sufficient-statistic components return 1 rather than
+        the true value.  Check the class attribute before relying on the
+        result for those components.
+
         Args:
             t: Frequencies in natural-parameter space, same structure as self.
 

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

@@ -31,7 +31,6 @@ def triangular_number_index(k: int, /) -> int | None:
 
 
 class Support:
-    @override
     def __init__(self, *, ring: Ring = real_field) -> None:
         super().__init__()
         self.ring = ring
@@ -280,6 +279,12 @@ class SquareMatrixSupport(Support):
         x = self.ring.unflattened(y, map_from_plane=map_from_plane)
         return xp.reshape(x, x.shape[:-1] + self.shape(dimensions))
 
+    @override
+    def generate(
+        self, xp: Namespace, rng: Generator, shape: Shape, safety: float, dimensions: int
+    ) -> JaxRealArray:
+        return self.ring.generate(xp, rng, (*shape, dimensions, dimensions), safety)
+
 
 class CircularBoundedSupport(VectorSupport):
     def __init__(self, radius: float) -> None:

{efax-2.2.0 → efax-2.2.4}/efax/_src/structure/assembler.py

@@ -68,7 +68,9 @@ class Assembler(Generic[P]):
 
         infos = []
         for info in self.infos:
-            assert issubclass(info.type_, ExpectationParametrization)
+            if not issubclass(info.type_, ExpectationParametrization):
+                msg = f"{info.type_.__name__} is not an EP"
+                raise TypeError(msg)
             infos.append(
                 SubDistributionInfo(
                     info.path,
@@ -85,7 +87,9 @@ class Assembler(Generic[P]):
 
         infos: list[SubDistributionInfo] = []
         for info in self.infos:
-            assert issubclass(info.type_, NaturalParametrization)
+            if not issubclass(info.type_, NaturalParametrization):
+                msg = f"{info.type_.__name__} is not an NP"
+                raise TypeError(msg)
             infos.append(
                 SubDistributionInfo(
                     info.path,
@@ -129,7 +133,7 @@ class Assembler(Generic[P]):
         ]
         q_values = parameters(q).values()
         q_params_as_p = dict(zip(p_paths, q_values, strict=True))
-        return self.assemble(q_params_as_p)
+        return self.assemble(q_params_as_p)  # ty: ignore
 
     def domain_support(self) -> dict[Path, Support]:
         """Return the domain support constraints for each simple sub-distribution in the tree."""

{efax-2.2.0 → efax-2.2.4}/efax/_src/structure/estimator.py

@@ -45,7 +45,9 @@ class Estimator(Assembler[P]):
             ExpectationParametrization,
         )
 
-        assert issubclass(type_p, ExpectationParametrization)
+        if not issubclass(type_p, ExpectationParametrization):
+            msg = f"{type_p.__name__} is not an EP"
+            raise TypeError(msg)
         return Estimator(
             [SubDistributionInfo((), type_p, 0, [])],
             {(name,): value for name, value in fixed_parameters.items()},
@@ -63,7 +65,9 @@ class Estimator(Assembler[P]):
         )
 
         infos = cls.create_assembler(p).infos
-        assert isinstance(p, ExpectationParametrization)
+        if not isinstance(p, ExpectationParametrization):
+            msg = f"{type(p).__name__} is not an EP"
+            raise TypeError(msg)
         fixed_parameters = parameters(p, fixed=True)
         return cls(infos, fixed_parameters)
 

{efax-2.2.0 → efax-2.2.4}/efax/_src/tools.py

@@ -4,7 +4,7 @@ from collections import defaultdict
 from collections.abc import Callable, Iterable, Mapping
 from functools import reduce
 from itertools import starmap
-from typing import TYPE_CHECKING, TypeVar
+from typing import TypeVar
 
 from array_api_compat import array_namespace
 from jax import jit
@@ -17,7 +17,7 @@ from .types import Axis
 
 
 @jit
-def parameter_dot_product(x: NaturalParametrization, y: Distribution, /) -> JaxRealArray:
+def parameter_dot_product(x: Distribution, y: Distribution, /) -> JaxRealArray:
     """Return the vectorized dot product over all of the variable parameters."""
 
     def dotted_fields() -> Iterable[JaxRealArray]:
@@ -38,12 +38,13 @@ T = TypeVar("T", bound=Distribution)
 
 
 def parameter_mean[T: Distribution](x: T, /, *, axis: Axis | None = None) -> T:
-    """Return the mean of the parameters (fixed and variable)."""
+    """Return the mean of the variable parameters; fixed parameters are preserved unchanged."""
     xp = array_namespace(x)
     structure = Assembler.create_assembler(x)
-    p = parameters(x)
-    q = {path: xp.mean(value, axis=axis) for path, value in p.items()}
-    return structure.assemble(q)
+    fixed = parameters(x, fixed=True)
+    variable = parameters(x, fixed=False)
+    averaged = {path: xp.mean(value, axis=axis) for path, value in variable.items()}
+    return structure.assemble({**fixed, **averaged})
 
 
 def parameter_map[T: Distribution](
@@ -80,7 +81,3 @@ def _parameter_dot_product(x: JaxComplexArray, y: JaxComplexArray, n_axes: int)
     xp = array_namespace(x, y)
     axes = tuple(range(-n_axes, 0))
     return xp.sum(xp.real(x * xp.conj(y)), axis=axes)
-
-
-if TYPE_CHECKING:
-    from .natural_parametrization import NaturalParametrization

{efax-2.2.0 → efax-2.2.4}/efax/_src/transform/joint.py

@@ -50,9 +50,19 @@ class JointDistribution(Distribution):
     @property
     @override
     def shape(self) -> Shape:
-        for distribution in self._sub_distributions.values():
-            return distribution.shape
-        raise ValueError
+        shapes = [d.shape for d in self._sub_distributions.values()]
+        if not shapes:
+            msg = "JointDistribution has no sub-distributions"
+            raise ValueError(msg)
+        shape = shapes[0]
+        if any(s != shape for s in shapes[1:]):
+            names = list(self._sub_distributions)
+            msg = (
+                f"Sub-distribution shapes are inconsistent: "
+                f"{dict(zip(names, shapes, strict=True))}"
+            )
+            raise ValueError(msg)
+        return shape
 
 
 @dataclass

{efax-2.2.0 → efax-2.2.4}/pyproject.toml

@@ -6,19 +6,18 @@ build-backend = "uv_build"
 dev = [
   "jupyter>=1",
   "lefthook>=2",
-  "pyright>=1.1.408",
   "pytest-xdist[psutil]>=3",
   "pytest>=9",
   "ruff>=0.15",
   "scipy-stubs>=1.15",
   "scipy>=1.15",
   "toml-sort>=0.24",
-  "ty>=0.0.22"
+  "ty>=0.0.29"
 ]
 
 [project]
 name = "efax"
-version = "2.2.0"
+version = "2.2.4"
 description = "Exponential families for JAX"
 readme = "README.rst"
 requires-python = ">=3.12, <3.15"
@@ -46,7 +45,7 @@ dependencies = [
   "opt-einsum>=3.4",
   "optimistix>=0.0.9",
   "optype[numpy]>=0.8.0",
-  "tjax>=1.4.4",
+  "tjax>=1.4.6",
   "typing_extensions>=4.8"
 ]
 

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

@@ -2,6 +2,7 @@
 
 from ._src.distributions.bernoulli import BernoulliEP, BernoulliNP
 from ._src.distributions.beta import BetaEP, BetaNP
+from ._src.distributions.categorical import CategoricalEP, CategoricalNP
 from ._src.distributions.chi import ChiEP, ChiNP
 from ._src.distributions.chi_square import ChiSquareEP, ChiSquareNP
 from ._src.distributions.cmvn.circularly_symmetric import (
@@ -31,7 +32,6 @@ from ._src.distributions.log_normal.unit_variance import (
     UnitVarianceLogNormalNP,
 )
 from ._src.distributions.logarithmic import LogarithmicEP, LogarithmicNP
-from ._src.distributions.categorical import CategoricalEP, CategoricalNP
 from ._src.distributions.multivariate_normal.arbitrary import (
     MultivariateNormalEP,
     MultivariateNormalNP,
@@ -105,6 +105,8 @@ __all__ = [
     "BetaEP",
     "BetaNP",
     "BooleanRing",
+    "CategoricalEP",
+    "CategoricalNP",
     "ChiEP",
     "ChiNP",
     "ChiSquareEP",
@@ -156,8 +158,6 @@ __all__ = [
     "LogarithmicEP",
     "LogarithmicNP",
     "Multidimensional",
-    "CategoricalEP",
-    "CategoricalNP",
     "MultivariateDiagonalNormalEP",
     "MultivariateDiagonalNormalNP",
     "MultivariateDiagonalNormalVP",