description: Annotates expected properties of a Tensor-valued distribution parameter.

tfp.substrates.jax.util.ParameterProperties

Annotates expected properties of a Tensor-valued distribution parameter.

Distributions and Bijectors implementing ._parameter_properties specify a ParameterProperties annotation for each of their Tensor-valued parameters.

Elements:

Batch shapes and parameter event_ndims

The batch_shape of a distribution/bijector/linear operator/PSD kernel/etc. instance is the shape of distinct parameterizations represented by that instance. It is computed by broadcasting the batch shapes of that instance's parameters, where an individual parameter's 'batch shape' is the shape of values specified for that parameter.

To compute the batch shape of a given parameter, we need to know what counts as a 'single value' of that parameter. For example, the scale parameter of tfd.Normal is semantically scalar-valued, so a value of shape [d] would have batch shape [d]. On the other hand, the scale_diag parameter of tfd.MultivariateNormalDiag is semantically vector-valued, so in this context a value of shaped [d] would have batch shape []. TFP formalizes this by annotating the scale parameter with event_ndims=0, and the scale_diag parameter with event_ndims=1.

In general, the event_ndims of a Tensor-valued parameter is the number of rightmost dimensions of its shape used to describe a single event of the parameterized instance. Equivalently, it is the minimal Tensor rank of a valid value for that parameter. The portion of each Tensor parameter's shape that remains after slicing off the rightmost event_ndims is its 'parameter batch shape'. The batch shape(s) of all parameters must broadcast with each other. For example, in a tfd.MultivariateNormalDiag(loc, scale_diag) distribution, where loc.shape == [3, 2] and scale_diag.shape == [4, 1, 2], the parameter batch shapes are [3] and [4, 1] respectively, and these broadcast to an overall batch shape of [4, 3].

Instance-dependent (callable) event_ndims

A parameter's event_ndims may be specified as a callable that returns an integer and takes as its argument an instance self of the class being parameterized. This allows parameters whose interpretation depends on other parameters. Callables for Bijector parameters must also accept a second argument x_event_ndims, described below.

For example, for the distribution parameter of tfd.Independent, we would specify event_ndims=lambda self: self.reinterpreted_batch_ndims, indicating that the outer class's relationship to the inner distribution depends on another instance parameter (reinterpreted_batch_ndims). The value returned from an event_ndims callable may be a Python int or an integer Tensor, but the callable itself may not cause graph side effects (e.g., create new Tensors). In cases where graph ops can't be avoided, the event_ndims callable should return None, and a separate callable event_ndims_tensor must be provided.

Parameters of non-Distribution objects

The notion of an 'event' generalizes beyond distributions. In general, an event refers to an instance of an object with batch_shape==[], and the event_ndims of a parameter describes the parameter value that would define such an instance. For example:

Non-Tensor-valued parameters (Distributions, Bijectors, etc).

The previous section discussed annotating parameters of non-Distribution objects. We'll now consider the orthogonal generalization: parameters that themselves take non-Tensor values. For example, the distribution and bijector parameters of tfd.TransformedDistribution are themselves a distribution and a bijector, respectively.

Structured parameter event_ndims

A parameter's event_ndims will be a nested structure of integers (list, dict, etc.) if either of the following applies:

  1. The parameter value itself is a nested structure. For example, in the joint bijector tfb.JointMap(bijectors=[tfb.Softplus(), tfb.Exp()]), the event_ndims of the bijectors parameter would be [0, 0], matching the structure of the bijectors value (note that since this structure is instance-dependent, the event_ndims would need to be specified using a callable, as detailed above).

  2. The parameter is a Bijector with structured forward_min_event_ndims. For example, in tfb.JointMap(bijectors=[tfb.Softplus(), tfb.Invert(tfb.Split(2))]), the event_ndims of the bijectors parameter would be [0, [1, 1]], since the inverse of the Split bijector has forward_min_event_ndims=[1, 1].

Any ambiguity between these two uses for structured event_ndims can be resolved by examining the parameter value. For example, event_ndims = [[2, 1], [1, 0]] could describe a nested structure containing four Tensors (or distributions, single-part bijectors, etc.), a list containing two structured bijectors, or a single bijector operating on nested lists of Tensors, but we can always tell which of these is the case by examining the actually instantiated parameter.

Note that JointDistribution-valued parameters never have structured event_ndims, despite having structured event shapes, because the event_ndims annotation of a Distribution parameter describes the number of that distribution's batch dimensions that contribute to an event of the outer parameterized object. Bijectors require additional annotation not because they operate on structured events, but rather because they operate in a context-specific manner depending on the event being transformed.

Choice of constraining bijectors

The practical support of a parameter---defined as the regime in which the distribution may be expected to produce numerically valid samples and (log-)densities---may differ slightly from the mathematical support. For example, Normal scale is mathematically supported on positive real numbers, but in practice, dividing by very small scales may cause overflow. We might therefore prefer a bijector such as tfb.Softplus(low=eps) that excludes very small values.

In general, default constraining bijectors should attempt to implement a practical rather than mathematical support, and users of default bijectors should be aware that extreme elements of the mathematical support may not be attainable. The notion of 'practical support' is inherently fuzzy, and defining it may require arbitrary choices. However, this is preferred to the alternative of allowing the default behavior to be numerically unstable in common settings. As a general guide, any restrictions on the mathematical support should be 'conceptually infinitesimal': it may be appropriate to constrain a Beta concentration parameter to be greater than eps, but not to be greater than 1 + eps, since the latter is a non-infinitesimal restriction of the mathematical support.

event_ndims

event_ndims_tensor

shape_fn

default_constraining_bijector_fn

is_preferred

is_tensor

specifies_shape

Methods

bijector_instance_event_ndims

View source

Computes parameter event_ndims when parameterizing a bijector.

instance_event_ndims

View source