efax 2.1.2__tar.gz → 2.2.3__tar.gz

{efax-2.1.2 → efax-2.2.3}/PKG-INFO

@@ -1,25 +1,21 @@
 Metadata-Version: 2.4
 Name: efax
-Version: 2.1.2
+Version: 2.2.3
 Summary: Exponential families for JAX
-Project-URL: source, https://github.com/NeilGirdhar/efax
+Author: Neil Girdhar
 Author-email: Neil Girdhar <mistersheik@gmail.com>
-Maintainer-email: Neil Girdhar <mistersheik@gmail.com>
 License-Expression: Apache-2.0
-License-File: LICENSE
 Classifier: Development Status :: 5 - Production/Stable
 Classifier: Intended Audience :: Science/Research
-Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Operating System :: OS Independent
-Classifier: Programming Language :: Python
 Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
 Classifier: Programming Language :: Python :: 3.14
+Classifier: Programming Language :: Python
 Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Classifier: Topic :: Software Development :: Libraries :: Python Modules
 Classifier: Typing :: Typed
-Requires-Python: <3.15,>=3.12
 Requires-Dist: array-api-compat>=1.10
 Requires-Dist: array-api-extra>=0.9.2
 Requires-Dist: jax>=0.9.0
@@ -27,8 +23,12 @@ 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>
+Requires-Python: >=3.12, <3.15
+Project-URL: source, https://github.com/NeilGirdhar/efax
 Description-Content-Type: text/x-rst
 
 .. role:: bash(code)
@@ -37,6 +37,8 @@ Description-Content-Type: text/x-rst
 .. role:: python(code)
     :language: python
 
+.. default-role:: python
+
 .. image:: https://img.shields.io/pypi/v/efax
    :target: https://pypi.org/project/efax/
    :alt: PyPI - Version
@@ -48,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
@@ -81,7 +89,7 @@ Framework
 =========
 Representation
 --------------
-EFAX has a single base class for its objects: :python:`Distribution` whose type encodes the
+EFAX has a single base class for its objects: `Distribution` whose type encodes the
 distribution family.
 
 Each parametrization object has a shape, and so it can store any number of distributions.
@@ -89,11 +97,11 @@ Operations on these objects are vectorized.
 This is unlike SciPy where each distribution is represented by a single object, and so a thousand
 distributions need a thousand objects, and corresponding calls to functions that operate on them.
 
-All parametrization objects are dataclasses using :python:`tjax.dataclass`.  These dataclasses are
-a modification of Python's dataclasses to support JAX's “PyTree” type registration.
+All parametrization objects are dataclasses using `tjax.dataclass`.  These dataclasses are
+a modification of Python's dataclasses to support JAX's "PyTree" type registration.
 
 Each of the fields of a parametrization object stores a parameter over a specified support.
-Some parameters are marked as “fixed”, which means that they are fixed with respect to the
+Some parameters are marked as "fixed", which means that they are fixed with respect to the
 exponential family.  An example of a fixed parameter is the failure number of the negative binomial
 distribution.
 
@@ -107,17 +115,17 @@ For example:
         negative_half_precision: RealArray = distribution_parameter(SymmetricMatrixSupport())
 
 In this case, we see that there are two natural parameters for the multivariate normal distribution.
-Objects of this type can hold any number of distributions:  if such an object :python:`x` has shape
-:python:`s`, then the shape of
-:python:`x.mean_times_precision` is :python:`(*s, n)` and the shape of
-:python:`x.negative_half_precision` is :python:`(*s, n, n)`.
+Objects of this type can hold any number of distributions:  if such an object `x` has shape
+`s`, then the shape of
+`x.mean_times_precision` is `(*s, n)` and the shape of
+`x.negative_half_precision` is `(*s, n, n)`.
 
 Parametrizations
 ----------------
 Each exponential family distribution has two special parametrizations: the natural and the
 expectation parametrization.  (These are described in the overview pdf.)
 Consequently, every distribution has at least two base classes, one inheriting from
-:python:`NaturalParametrization` and one from :python:`ExpectationParametrization`.
+`NaturalParametrization` and one from `ExpectationParametrization`.
 
 The motivation for the natural parametrization is combining and scaling independent predictive
 evidence.  In the natural parametrization, these operations correspond to scaling and addition.
@@ -127,56 +135,129 @@ maximum likelihood distribution that could have produced them.  In the expectati
 this is an expected value.
 
 EFAX provides conversions between the two parametrizations through the
-:python:`NaturalParametrization.to_exp` and :python:`ExpectationParametrization.to_nat` methods.
+`NaturalParametrization.to_exp` and `ExpectationParametrization.to_nat` methods.
+
+Some distributions also provide additional convenience parametrizations.  For example, the normal
+distribution offers `NormalVP` (variance parametrization) and `NormalDP` (deviation
+parametrization), and the multivariate normal provides `MultivariateNormalVP` and
+`MultivariateDiagonalNormalVP`.  These are not exponential-family parametrizations but are
+useful for constructing and interpreting distributions.
 
 Important methods
 -----------------
 EFAX aims to provide the main methods used in machine learning.
 
-Every :python:`Distribution` has methods:
+Every `Distribution` has:
 
-- :python:`flattened` and :python:`unflattened` to flatten and unflatten the parameters into a
-  single array.  Typically, array-valued signals in a machine learning model would be unflattened
-  into a distribution object, operated on, and then flattened before being sent back to the model.
-  Flattening is careful with distributions with symmetric (or Hermitian) matrix-valued parameters.
-  It only stores the upper triangular elements.  And,
-- :python:`shape`, which supports broadcasting.
+- `shape` and `ndim`, which support broadcasting, and
+- indexing via `[]`, which slices all parameter arrays simultaneously.
 
-Every :python:`NaturalParametrization` has methods:
+Every `NaturalParametrization` has methods:
 
-- :python:`to_exp` to convert itself to expectation parameters.
-- :python:`sufficient_statistics` to produce the sufficient statistics given an observation (used in
+- `to_exp` to convert itself to expectation parameters,
+- `sufficient_statistics` to produce the sufficient statistics given an observation (used in
   maximum likelihood estimation),
-- :python:`pdf` and :python:`log_pdf`, which is the density or mass function and its logarithm,
-- :python:`fisher_information`, which is the Fisher information matrix, and
-- :python:`kl_divergence`, which is the KL divergence.
-
-Every :python:`ExpectationParametrization` has methods:
-
-- :python:`to_nat` to convert itself to natural parameters, and
-- :python:`kl_divergence`, which is the KL divergence.
+- `log_normalizer`, the log partition function,
+- `carrier_measure`, the base measure,
+- `pdf` and `log_pdf`, which are the density or mass function and its logarithm,
+- `fisher_information_diagonal` and `fisher_information_trace`, which return the
+  diagonal and trace of the Fisher information matrix stored as distribution objects,
+- `apply_fisher_information`, which applies the Fisher information matrix to a vector of
+  expectation parameters efficiently in a single VJP pass,
+- `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 (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:
+
+- `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:
 
-- :python:`HasConjugatePrior` can produce the conjugate prior,
-- :python:`HasGeneralizedConjugatePrior` can produce a generalization of the conjugate prior,
-- :python:`Multidimensional` distributions have a integer number of `dimensions`, and
-- :python:`Samplable` distributions support sampling.
+- `HasConjugatePrior` can produce and recover the conjugate prior,
+- `HasGeneralizedConjugatePrior` extends that with per-dimension pseudo-observation counts,
+- `Multidimensional` distributions have an integer number of `dimensions`, and
+- `Samplable` distributions support sampling.
 
 Some parametrizations inherit from these public mixins:
 
-- :python:`HasEntropyEP` is an expectation parametrization with an entropy and cross entropy, and
-- :python:`HasEntropyNP` is a natural parametrization with an entropy,  (The cross entropy is not
-  efficient.)
+- `HasEntropy` is a distribution with a `entropy` method,
+- `HasEntropyEP` is an expectation parametrization with analytically tractable entropy and
+  `cross_entropy`, and
+- `HasEntropyNP` is a natural parametrization with analytically tractable entropy via the
+  paired expectation parametrization.
 
 Some parametrizations inherit from these private mixins:
 
-- :python:`ExpToNat` implements the conversion from expectation to natural parameters when no
+- `ExpToNat` implements the conversion from expectation to natural parameters when no
   analytical solution is possible.  It uses Newton's method with a Jacobian to invert the gradient
   log-normalizer.
-- :python:`TransformedNaturalParametrization` produces a natural parametrization by relating it to
+- `TransformedNaturalParametrization` produces a natural parametrization by relating it to
   an existing natural parametrization.  And similarly for
-  :python:`TransformedExpectationParametrization`.
+  `TransformedExpectationParametrization`.
+
+Joint distributions
+-------------------
+`JointDistribution`, `JointDistributionE`, and `JointDistributionN` compose
+multiple independent distributions into a single object.  `JointDistributionE` holds
+expectation parametrizations and implements `HasEntropyEP`; `JointDistributionN`
+holds natural parametrizations.  They support the same `to_nat` / `to_exp`
+conversions as simple distributions.
+
+Structure utilities
+-------------------
+EFAX provides three classes that capture the static metadata of a distribution tree—its types,
+parameter names, and dimension information—without requiring a live instance.  They form an
+inheritance hierarchy:
+
+**Assembler**
+  Stores a post-order traversal of a `Distribution` tree (types, paths, dimensions) so
+  that distributions can be reconstructed from raw parameter data without passing type information
+  alongside arrays.  Key methods:
+
+  - `assemble(params)` — rebuild a `Distribution` from a `{path: array}` mapping,
+  - `coerce_from_distribution(q)` — reinterpret q's numeric values under this Assembler's types,
+  - `domain_support()` — enumerate each leaf distribution's parameter constraints,
+  - `generate_random(xp, rng, shape, safety)` — draw a random distribution with valid parameters, and
+  - `to_nat()` / `to_exp()` — return a copy whose types are all in natural or expectation form.
+
+**Estimator** *(extends Assembler)*
+  Adds maximum likelihood estimation by recording which parameters are fixed (held constant across
+  observations) and which are free.  Because the MLE for every exponential family equals the mean
+  of the sufficient statistics, estimation reduces to a single call:
+
+  - `sufficient_statistics(x)` — compute the sufficient statistics of observation x, with
+    fixed parameters supplied automatically.
+
+  Create one with `Estimator.from_type(type_p, **fixed)`,
+  `Estimator.from_expectation(p)`, or `Estimator.from_natural(p)`.
+
+**Flattener** *(extends Estimator)*
+  Adds encoding and decoding between a `Distribution` and an array of shape
+  `(*distribution.shape, k)`, making distributions compatible with neural networks and numerical
+  optimizers.  Fixed parameters are excluded from the encoded array and reinserted automatically
+  on decode.
+
+  - `Flattener.flatten(p, mapped_to_plane=True)` — encode p into a `(Flattener, array)` pair,
+  - `unflatten(array)` — decode the array back into a distribution, and
+  - `final_dimension_size()` — the size k of the last axis of the encoded array.
+
+  The `mapped_to_plane` flag controls whether constrained parameters (e.g., those on a
+  simplex or restricted to the positive reals) are bijectively mapped to all of ℝⁿ.  Set it
+  `True` when passing to a neural network (to prevent invalid outputs), and `False`
+  when the raw magnitudes matter—for example when differencing expectation parameters or computing
+  Jacobians.
 
 Distributions
 =============
@@ -220,7 +301,7 @@ EFAX supports the following distributions:
 - on a finite set:
 
   - Bernoulli
-  - multinomial
+  - categorical
 
 - on the nonnegative integers:
 
@@ -246,9 +327,14 @@ EFAX supports the following distributions:
   - Dirichlet
   - generalized Dirichlet
 
-- on the n-sphere:
+- on spheres:
 
   - von Mises-Fisher
+  - complex von Mises
+
+- on positive-definite matrices:
+
+  - Wishart
 
 Usage
 =====
@@ -376,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]))
@@ -395,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.
@@ -437,14 +532,14 @@ instead.
     likelihood estimation.
 
     Suppose you have some samples from a distribution family with unknown
-    parameters, and you want to estimate the maximum likelihood parmaters of the
+    parameters, and you want to estimate the maximum likelihood parameters of the
     distribution.
     """
     import jax.numpy as jnp
     import jax.random as jr
     from tjax import print_generic
 
-    from efax import DirichletEP, DirichletNP, MaximumLikelihoodEstimator, parameter_mean
+    from efax import DirichletEP, DirichletNP, Estimator, parameter_mean
 
     # Consider a Dirichlet distribution with a given alpha.
     alpha = jnp.asarray([2.0, 3.0, 4.0])
@@ -457,12 +552,14 @@ instead.
 
     # Now, let's find the maximum likelihood Dirichlet distribution that fits it.
     # First, convert the samples to their sufficient statistics.
-    estimator = MaximumLikelihoodEstimator.create_simple_estimator(DirichletEP)
+    estimator = Estimator.from_type(DirichletEP)
     ss = estimator.sufficient_statistics(samples)
-    # ss has type DirichletEP.  This is similar to the conjguate prior of the
+    # ss has type DirichletEP.  This is similar to the conjugate prior of the
     # 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.