An example of a fixed parameter is the failure number of the negative binomialįor example: class MultivariateNormalNP ( NaturalParametrization ): mean_times_precision : RealArray = distribution_parameter ( VectorSupport ()) negative_half_precision : RealArray = distribution_parameter ( SymmetricMatrixSupport ())
Some parameters are marked as “fixed”, which means that they are fixed with respect to theĮxponential family. These dataclasses areĪ modification of Python’s dataclasses to support JAX’s “PyTree” type registration.Įach of the fields of a parametrization object stores a parameter over a specified support. This is unlike SciPy where each distribution is represented by a single object, and so a thousandĭistributions need a thousand objects, and corresponding calls to functions that operate on them.Īll parametrization objects are dataclasses using tjax. Operations on these objects are vectorized.
Framework RepresentationĮFAX has a single base class for its objects: Parametrization whose type encodes theĮach parametrization object has a shape, and so it can store any number of distributions. Way to implement cross entropy between X and Y relies on X being in the expectation parametrizationĪnd Y in the natural parametrization. An example of why this matters is that the most efficient Natural and expectation parametrizations-and a uniform interface to efficient implementations of the In JAX is that EFAX provides the two most important parametrizations for each exponential family-the The main motivation for using EFAX over a library like tensorflow-probability or the basic functions Normal, gamma, beta, exponential, Poisson, binomial, and Bernoulli distributions.įor an explanation of the fundamental ideas behind this library, see our overview on exponential The exponential families are an important class of probability distributions that include the This library provides a set of tools for working with exponential family distributions in the