def __init__(self, count_distribution, inflated_distribution=None, logits=None, probs=None, eps=1e-8, validate_args=False, allow_nan_stats=True, name="ZeroInflated"): r"""Initialize a zero-inflated distribution. A `ZeroInflated` is defined by a zero-inflation rate (`inflated_distribution`, representing the probabilities of excess zeros) and a `Distribution` object having matching dtype, batch shape, event shape, and continuity properties (the dist). Arguments: count_distribution : A `tfp.distributions.Distribution` instance. The instance must have `batch_shape` matching the zero-inflation distribution. inflated_distribution: `tfp.distributions.Bernoulli`-like instance. Manages the probability of excess zeros, the zero-inflated rate. Must have either scalar `batch_shape` or `batch_shape` matching `count_distribution.batch_shape`. logits: An N-D `Tensor` representing the log-odds of a excess zeros A zero-inflation rate, where the probability of excess zeros is sigmoid(logits). Only one of `logits` or `probs` should be passed in. probs: An N-D `Tensor` representing the probability of a zero event. Each entry in the `Tensor` parameterizes an independent ZeroInflated distribution. Only one of `logits` or `probs` should be passed in. validate_args: Python `bool`, default `False`. If `True`, raise a runtime error if batch or event ranks are inconsistent between pi and any of the distributions. This is only checked if the ranks cannot be determined statically at graph construction time. allow_nan_stats: Boolean, default `True`. If `False`, raise an exception if a statistic (e.g. mean/mode/etc...) is undefined for any batch member. If `True`, batch members with valid parameters leading to undefined statistics will return NaN for this statistic. name: A name for this distribution (optional). References: Liu, L. & Blei, D.M.. (2017). Zero-Inflated Exponential Family Embeddings. Proceedings of the 34th International Conference on Machine Learning, in PMLR 70:2140-2148 """ parameters = dict(locals()) with tf.compat.v1.name_scope(name) as name: # main count distribution if not isinstance(count_distribution, distribution.Distribution): raise TypeError( "count_distribution must be a Distribution instance" " but saw: %s" % count_distribution) # Zero inflation distribution if inflated_distribution is None: inflated_distribution = Bernoulli( logits=logits, probs=probs, dtype=tf.int32, validate_args=validate_args, allow_nan_stats=allow_nan_stats, name="ZeroInflatedRate") elif not isinstance(inflated_distribution, distribution.Distribution): raise TypeError( "inflated_distribution must be a Distribution instance" " but saw: %s" % inflated_distribution) # Matching the event shape inflated_ndim = len(inflated_distribution.event_shape) count_ndim = len(count_distribution.event_shape) if inflated_ndim < count_ndim: inflated_distribution = Independent(inflated_distribution, count_ndim - inflated_ndim) self._count_distribution = count_distribution self._inflated_distribution = inflated_distribution # if self._count_distribution.batch_shape.ndims is None: raise ValueError( "Expected to know rank(batch_shape) from count_distribution" ) if self._inflated_distribution.batch_shape.ndims is None: raise ValueError( "Expected to know rank(batch_shape) from inflated_distribution" ) self._eps = tensor_util.convert_nonref_to_tensor( eps, dtype_hint=count_distribution.dtype, name='eps') # We let the zero-inflated distribution access _graph_parents since its arguably # more like a baseclass. super(ZeroInflated, self).__init__( dtype=self._count_distribution.dtype, reparameterization_type=self._count_distribution. reparameterization_type, validate_args=validate_args, allow_nan_stats=allow_nan_stats, parameters=parameters, name=name, )
def __init__(self, transform_fn, pretransformed_input, dtype=None, shape=NONE_SPECIFIED, name=None): """Creates the `DeferredTensor` object. Args: transform_fn: Python `callable` or `tfp.bijectors.Bijector`-like instance. When `callable`, should take `pretransformed_input` and return a `Tensor` (representing by this object). pretransformed_input: object with `shape`, `dtype` properties (typically a `tf.Variable`) passed into `transform_fn` when this object is acted upon in a `Tensor` context, eg, `tf.convert_to_tensor`, `+`, `tf.math.exp`, etc. dtype: Equivalent to what would otherwise be `transform_fn(pretransformed_input).dtype`. Default value: `None` (i.e., `getattr(transform_fn, 'dtype', None) or pretransformed_input.dtype`). shape: Equivalent to what would otherwise be `transform_fn(pretransformed_input).shape`. Default value: `'None'` (i.e., `getattr(transform_fn, 'forward_event_shape', lambda x: x)( pretransformed_input.shape)`). name: Python `str` representing this object's `name`; used only in graph mode. Default value: `None` (i.e., `(getattr(transform_fn, 'name', None) or transform_fn.__name__ + '_' + pretransformed_input.name)`). Raises: TypeError: if `transform_fn` is not `callable`. TypeError: if `pretransformed_input` lacks `dtype` and/or `shape` properties (and `dtype` and/or `shape` arguments are unspecified). """ pretransformed_input = tensor_util.convert_nonref_to_tensor( pretransformed_input, name='pretransformed_input') if dtype is None: dtype = (getattr(transform_fn, 'dtype', None) or dtype_util.base_dtype(pretransformed_input.dtype)) try: dtype = None if dtype is None else tf.dtypes.as_dtype(dtype) except TypeError: raise TypeError('Argument `dtype` must be convertible to a ' '`tf.dtypes.DType`; saw "{}" of type "{}".'.format( repr(dtype), type(dtype))) if shape == NONE_SPECIFIED: shape = getattr(transform_fn, 'forward_event_shape', _identity) shape = shape(pretransformed_input.shape) try: shape = tf.TensorShape(shape) except TypeError: raise TypeError('Argument `shape` must be convertible to a ' '`tf.TensorShape`; saw "{}".'.format(shape)) name = name or getattr(transform_fn, 'name', None) if not name: name = '_'.join([ transform_fn.__name__, getattr(pretransformed_input, 'name', '')]) name = name_util.strip_invalid_chars(name) name = name_util.camel_to_lower_snake(name) name = name_util.get_name_scope_name(name) name = name_util.strip_invalid_chars(name) if hasattr(transform_fn, 'forward'): fwd_name = '"{}"'.format(transform_fn.name) transform_fn = transform_fn.forward else: fwd_name = transform_fn.__name__ if not callable(transform_fn): raise TypeError('Argument `transform_fn` must be `callable`.') super(DeferredTensor, self).__init__(name=name) self._transform_fn = transform_fn self._pretransformed_input = pretransformed_input self._dtype = dtype self._shape = shape self._name = name self._fwd_name = fwd_name # Secret handshake with tf.is_tensor to return True for DT. # # Works around an exception in LinearOperator (which in 2.0.0 checks only # `tf.is_tensor`, not also `linear_operator_util.is_ref`: # ValueError: Graph parent item 0 is not a Tensor; # <DeferredTensor: dtype=float32, shape=[2], fn=exp>. # TODO(b/140157055): Remove this shim after LinOp is patched in 2.0. self.is_tensor_like = True
def __init__(self, rate1=None, rate2=None, log_rate1=None, log_rate2=None, force_probs_to_zero_outside_support=False, validate_args=False, allow_nan_stats=True, name='Skellam'): """Initialize a batch of Skellam distributions. Args: rate1: Floating point tensor, the first rate parameter. `rate1` must be positive. Must specify exactly one of `rate1` and `log_rate1` rate2: Floating point tensor, the second rate parameter. `rate` must be positive. Must specify exactly one of `rate2` and `log_rate2`. log_rate1: Floating point tensor, the log of the first rate parameter. Must specify exactly one of `rate1` and `log_rate1`. log_rate2: Floating point tensor, the log of the second rate parameter. Must specify exactly one of `rate2` and `log_rate2`. force_probs_to_zero_outside_support: Python `bool`. When `True`, `log_prob` returns `-inf` (and `prob` returns `0`) for non-integer inputs. When `False`, `log_prob` evaluates the Skellam pmf as a continuous function (note that this function is not itself a normalized probability log-density). Default value: `False`. validate_args: Python `bool`. When `True` distribution parameters are checked for validity despite possibly degrading runtime performance. When `False` invalid inputs may silently render incorrect outputs. Default value: `False`. allow_nan_stats: Python `bool`. When `True`, statistics (e.g., mean, mode, variance) use the value "`NaN`" to indicate the result is undefined. When `False`, an exception is raised if one or more of the statistic's batch members are undefined. Default value: `True`. name: Python `str` name prefixed to Ops created by this class. Raises: ValueError: if none or both of `rate1`, `log_rate1` are specified. ValueError: if none or both of `rate2`, `log_rate2` are specified. """ parameters = dict(locals()) if (rate1 is None) == (log_rate1 is None): raise ValueError( 'Must specify exactly one of `rate1` and `log_rate1`.') if (rate2 is None) == (log_rate2 is None): raise ValueError( 'Must specify exactly one of `rate2` and `log_rate2`.') with tf.name_scope(name) as name: dtype = dtype_util.common_dtype( [rate1, rate2, log_rate1, log_rate2], dtype_hint=tf.float32) self._rate1 = tensor_util.convert_nonref_to_tensor(rate1, name='rate1', dtype=dtype) self._log_rate1 = tensor_util.convert_nonref_to_tensor( log_rate1, name='log_rate1', dtype=dtype) self._rate2 = tensor_util.convert_nonref_to_tensor(rate2, name='rate2', dtype=dtype) self._log_rate2 = tensor_util.convert_nonref_to_tensor( log_rate2, name='log_rate2', dtype=dtype) self._force_probs_to_zero_outside_support = force_probs_to_zero_outside_support super(Skellam, self).__init__( dtype=dtype, reparameterization_type=reparameterization.NOT_REPARAMETERIZED, validate_args=validate_args, allow_nan_stats=allow_nan_stats, parameters=parameters, name=name)
def test_np_object(self): x = np.array(0.) y = tensor_util.convert_nonref_to_tensor(x) self.assertIsInstance(y, tf.Tensor) self.assertEqual(x, self.evaluate(y))
def __init__(self, loc, scale_row, scale_column, validate_args=False, allow_nan_stats=True, name='MatrixNormalLinearOperator'): """Construct Matrix Normal distribution on `R^{n x p}`. The `batch_shape` is the broadcast shape between `loc`, `scale_row` and `scale_column` arguments. The `event_shape` is given by the matrix implied by `loc`. Args: loc: Floating-point `Tensor`, having shape `[B1, ..., Bb, n, p]`. scale_row: Instance of `LinearOperator` with the same `dtype` as `loc` and shape `[B1, ..., Bb, n, n]`. scale_column: Instance of `LinearOperator` with the same `dtype` as `loc` and shape `[B1, ..., Bb, p, p]`. validate_args: Python `bool`, default `False`. Whether to validate input with asserts. If `validate_args` is `False`, and the inputs are invalid, correct behavior is not guaranteed. allow_nan_stats: Python `bool`, default `True`. If `False`, raise an exception if a statistic (e.g. mean/mode/etc...) is undefined for any batch member If `True`, batch members with valid parameters leading to undefined statistics will return NaN for this statistic. name: The name to give Ops created by the initializer. """ parameters = dict(locals()) with tf.name_scope(name) as name: dtype = dtype_util.common_dtype([loc, scale_column, scale_row], dtype_hint=tf.float32) loc = tensor_util.convert_nonref_to_tensor(loc, dtype=dtype, name='loc') self._loc = loc if not hasattr(scale_row, 'matmul'): raise ValueError( '`scale_row` must be a `tf.linalg.LinearOperator`.') if not hasattr(scale_column, 'matmul'): raise ValueError( '`scale_column` must be a `tf.linalg.LinearOperator`.') if validate_args and not scale_row.is_non_singular: raise ValueError('`scale_row` must be non-singular.') if validate_args and not scale_column.is_non_singular: raise ValueError('`scale_column` must be non-singular.') self._scale_row = scale_row self._scale_column = scale_column super(MatrixNormalLinearOperator, self).__init__( dtype=dtype, validate_args=validate_args, allow_nan_stats=allow_nan_stats, reparameterization_type=reparameterization.FULLY_REPARAMETERIZED, parameters=parameters, name=name) self._parameters = parameters
def __init__(self, loc, scale, low, high, validate_args=False, allow_nan_stats=True, name='TruncatedNormal'): """Construct TruncatedNormal. All parameters of the distribution will be broadcast to the same shape, so the resulting distribution will have a batch_shape of the broadcast shape of all parameters. Args: loc: Floating point tensor; the mean of the normal distribution(s) ( note that the mean of the resulting distribution will be different since it is modified by the bounds). scale: Floating point tensor; the std deviation of the normal distribution(s). low: `float` `Tensor` representing lower bound of the distribution's support. Must be such that `low < high`. high: `float` `Tensor` representing upper bound of the distribution's support. Must be such that `low < high`. validate_args: Python `bool`, default `False`. When `True` distribution parameters are checked at run-time. allow_nan_stats: Python `bool`, default `True`. When `True`, statistics (e.g., mean, mode, variance) use the value '`NaN`' to indicate the result is undefined. When `False`, an exception is raised if one or more of the statistic's batch members are undefined. name: Python `str` name prefixed to Ops created by this class. """ parameters = dict(locals()) with tf.name_scope(name) as name: dtype = dtype_util.common_dtype([loc, scale, low, high], tf.float32) self._loc = tensor_util.convert_nonref_to_tensor(loc, name='loc', dtype=dtype) self._scale = tensor_util.convert_nonref_to_tensor(scale, name='scale', dtype=dtype) self._low = tensor_util.convert_nonref_to_tensor(low, name='low', dtype=dtype) self._high = tensor_util.convert_nonref_to_tensor(high, name='high', dtype=dtype) dtype_util.assert_same_float_dtype( [self._loc, self._scale, self._low, self._high]) super(TruncatedNormal, self).__init__( dtype=dtype, # This distribution is fully reparameterized. loc, scale have straight # through gradients. The gradients for the bounds are implemented # using custom derived expressions based on implicit gradients. # For the special case of lower bound zero and a positive upper bound # an equivalent expression can also be found in Sec 9.1.1. # of https://arxiv.org/pdf/1806.01851.pdf. The implementation here # handles arbitrary bounds. reparameterization_type=reparameterization. FULLY_REPARAMETERIZED, validate_args=validate_args, allow_nan_stats=allow_nan_stats, parameters=parameters, name=name)
def __init__(self, low=None, high=None, hinge_softness=None, validate_args=False, name='soft_clip'): """Instantiates the SoftClip bijector. Args: low: Optional float `Tensor` lower bound. If `None`, the lower-bound constraint is omitted. Default value: `None`. high: Optional float `Tensor` upper bound. If `None`, the upper-bound constraint is omitted. Default value: `None`. hinge_softness: Optional nonzero float `Tensor`. Controls the softness of the constraint at the boundaries; values outside of the constraint set are mapped into intervals of width approximately `log(2) * hinge_softness` on the interior of each boundary. High softness reserves more space for values outside of the constraint set, leading to greater distortion of inputs *within* the constraint set, but improved numerical stability near the boundaries. Default value: `None` (`1.0`). validate_args: Python `bool` indicating whether arguments should be checked for correctness. name: Python `str` name given to ops managed by this object. """ with tf.name_scope(name): dtype = dtype_util.common_dtype([low, high, hinge_softness], dtype_hint=tf.float32) low = tensor_util.convert_nonref_to_tensor(low, name='low', dtype=dtype) high = tensor_util.convert_nonref_to_tensor(high, name='high', dtype=dtype) hinge_softness = tensor_util.convert_nonref_to_tensor( hinge_softness, name='hinge_softness', dtype=dtype) softplus_bijector = softplus.Softplus( hinge_softness=hinge_softness) negate = tf.convert_to_tensor(-1., dtype=dtype) components = [] if low is not None and high is not None: # Support reference tensors (eg Variables) for `high` and `low` by # deferring all computation on them until needed. width = tfp_util.DeferredTensor( pretransformed_input=high, transform_fn=lambda high: high - low) negated_shrinkage_factor = tfp_util.DeferredTensor( pretransformed_input=width, transform_fn=lambda w: tf.cast( # pylint: disable=g-long-lambda negate * w / softplus_bijector.forward(w), dtype=dtype)) # Implement the soft constraint from 'Mathematical Details' above: # softclip(x) := -softplus(width - softplus(x - low)) * # (width) / (softplus(width)) + high components = [ shift.Shift(high), scale.Scale(negated_shrinkage_factor), softplus_bijector, shift.Shift(width), scale.Scale(negate), softplus_bijector, shift.Shift(tfp_util.DeferredTensor(low, lambda x: -x)) ] elif low is not None: # Implement a soft lower bound: # softlower(x) := softplus(x - low) + low components = [ shift.Shift(low), softplus_bijector, shift.Shift(tfp_util.DeferredTensor(low, lambda x: -x)) ] elif high is not None: # Implement a soft upper bound: # softupper(x) := -softplus(high - x) + high components = [ shift.Shift(high), scale.Scale(negate), softplus_bijector, scale.Scale(negate), shift.Shift(high) ] self._low = low self._high = high self._hinge_softness = hinge_softness self._chain = chain.Chain(components, validate_args=validate_args) super(SoftClip, self).__init__(forward_min_event_ndims=0, dtype=dtype, validate_args=validate_args, is_constant_jacobian=not components, name=name)
def __init__(self, mean_direction, concentration, validate_args=False, allow_nan_stats=True, name='VonMisesFisher'): """Creates a new `VonMisesFisher` instance. Args: mean_direction: Floating-point `Tensor` with shape [B1, ... Bn, D]. A unit vector indicating the mode of the distribution, or the unit-normalized direction of the mean. (This is *not* in general the mean of the distribution; the mean is not generally in the support of the distribution.) NOTE: `D` is currently restricted to <= 5. concentration: Floating-point `Tensor` having batch shape [B1, ... Bn] broadcastable with `mean_direction`. The level of concentration of samples around the `mean_direction`. `concentration=0` indicates a uniform distribution over the unit hypersphere, and `concentration=+inf` indicates a `Deterministic` distribution (delta function) at `mean_direction`. validate_args: Python `bool`, default `False`. When `True` distribution parameters are checked for validity despite possibly degrading runtime performance. When `False` invalid inputs may silently render incorrect outputs. allow_nan_stats: Python `bool`, default `True`. When `True`, statistics (e.g., mean, mode, variance) use the value "`NaN`" to indicate the result is undefined. When `False`, an exception is raised if one or more of the statistic's batch members are undefined. name: Python `str` name prefixed to Ops created by this class. Raises: ValueError: For known-bad arguments, i.e. unsupported event dimension. """ parameters = dict(locals()) with tf.name_scope(name) as name: dtype = dtype_util.common_dtype([mean_direction, concentration], tf.float32) self._mean_direction = tensor_util.convert_nonref_to_tensor( mean_direction, name='mean_direction', dtype=dtype) self._concentration = tensor_util.convert_nonref_to_tensor( concentration, name='concentration', dtype=dtype) static_event_dim = tf.compat.dimension_value( tensorshape_util.with_rank_at_least(self._mean_direction.shape, 1)[-1]) if static_event_dim is not None and static_event_dim > 5: raise ValueError('von Mises-Fisher ndims > 5 is not currently ' 'supported') # mean_direction is always reparameterized. # concentration is only for event_dim==3, via an inversion sampler. reparameterization_type = (reparameterization.FULLY_REPARAMETERIZED if static_event_dim == 3 else reparameterization.NOT_REPARAMETERIZED) super(VonMisesFisher, self).__init__( dtype=self._concentration.dtype, validate_args=validate_args, allow_nan_stats=allow_nan_stats, reparameterization_type=reparameterization_type, parameters=parameters, name=name)
def __init__(self, initial_distribution, transition_distribution, observation_distribution, num_steps, validate_args=False, allow_nan_stats=True, time_varying_observation_distribution=False, name='HiddenMarkovModel'): """Initialize hidden Markov model. Args: initial_distribution: A `Categorical`-like instance. Determines probability of first hidden state in Markov chain. The number of categories must match the number of categories of `transition_distribution` as well as both the rightmost batch dimension of `transition_distribution` and the rightmost batch dimension of `observation_distribution`. transition_distribution: A `Categorical`-like instance. The rightmost batch dimension indexes the probability distribution of each hidden state conditioned on the previous hidden state. observation_distribution: A `tfp.distributions.Distribution`-like instance. The rightmost batch dimension indexes the distribution of each observation conditioned on the corresponding hidden state. num_steps: The number of steps taken in Markov chain. An integer valued tensor. The number of transitions is `num_steps - 1`. validate_args: Python `bool`, default `False`. When `True` distribution parameters are checked for validity despite possibly degrading runtime performance. When `False` invalid inputs may silently render incorrect outputs. Default value: `False`. allow_nan_stats: Python `bool`, default `True`. When `True`, statistics (e.g., mean, mode, variance) use the value "`NaN`" to indicate the result is undefined. When `False`, an exception is raised if one or more of the statistic's batch members are undefined. Default value: `True`. time_varying_observation_distribution: Python `bool`, default `False`. When `True`, the observation_distribution has an additional batch dimension that indexes the distribution of each observation conditioned on the corresponding timestep. This dimension size should always match num_steps and is the second-to-last batch axis in the batch dimensions (just to the left of the dimension for the number of states). name: Python `str` name prefixed to Ops created by this class. Default value: "HiddenMarkovModel". Raises: ValueError: if `num_steps` is not at least 1. ValueError: if `initial_distribution` does not have scalar `event_shape`. ValueError: if `transition_distribution` does not have scalar `event_shape.` ValueError: if `transition_distribution` and `observation_distribution` are fully defined but don't have matching rightmost dimension. """ parameters = dict(locals()) # pylint: disable=protected-access with tf.name_scope(name) as name: self._num_steps = tensor_util.convert_nonref_to_tensor(num_steps) self._initial_distribution = initial_distribution self._observation_distribution = observation_distribution self._transition_distribution = transition_distribution self._time_varying_observation_distribution = ( time_varying_observation_distribution) num_steps_ = tf.get_static_value(num_steps) if num_steps_ is not None: if np.ndim(num_steps_) != 0: raise ValueError( '`num_steps` must be a scalar but it has rank {}'.format( np.ndim(num_steps_))) else: self._static_event_shape = tf.TensorShape( [num_steps_]).concatenate( self._observation_distribution.event_shape) else: self._static_event_shape = tf.TensorShape( [None]).concatenate( self._observation_distribution.event_shape) observation_batch_shape = ( self._observation_distribution.batch_shape[:-2] if self._time_varying_observation_distribution else self._observation_distribution.batch_shape[:-1]) self._static_batch_shape = tf.broadcast_static_shape( self._initial_distribution.batch_shape, tf.broadcast_static_shape( self._transition_distribution.batch_shape[:-1], observation_batch_shape)) # pylint: disable=protected-access super(HiddenMarkovModel, self).__init__( dtype=self._observation_distribution.dtype, reparameterization_type=reparameterization.NOT_REPARAMETERIZED, validate_args=validate_args, allow_nan_stats=allow_nan_stats, parameters=parameters, name=name) # pylint: enable=protected-access self._parameters = parameters
def __init__(self, distribution, bijector, batch_shape=None, event_shape=None, kwargs_split_fn=_default_kwargs_split_fn, validate_args=False, parameters=None, name=None): """Construct a Transformed Distribution. Args: distribution: The base distribution instance to transform. Typically an instance of `Distribution`. bijector: The object responsible for calculating the transformation. Typically an instance of `Bijector`. batch_shape: `integer` vector `Tensor` which overrides `distribution` `batch_shape`; valid only if `distribution.is_scalar_batch()`. event_shape: `integer` vector `Tensor` which overrides `distribution` `event_shape`; valid only if `distribution.is_scalar_event()`. kwargs_split_fn: Python `callable` which takes a kwargs `dict` and returns a tuple of kwargs `dict`s for each of the `distribution` and `bijector` parameters respectively. Default value: `_default_kwargs_split_fn` (i.e., `lambda kwargs: (kwargs.get('distribution_kwargs', {}), kwargs.get('bijector_kwargs', {}))`) validate_args: Python `bool`, default `False`. When `True` distribution parameters are checked for validity despite possibly degrading runtime performance. When `False` invalid inputs may silently render incorrect outputs. parameters: Locals dict captured by subclass constructor, to be used for copy/slice re-instantiation operations. name: Python `str` name prefixed to Ops created by this class. Default: `bijector.name + distribution.name`. Raises: ValueError: If `distribution` is a joint distribution and a `batch_shape` override is passed. ValueError: If `distribution` is a joint distribution and an `event_shape` override is passed. """ parameters = dict(locals()) if parameters is None else parameters name = name or (('' if bijector is None else bijector.name) + (distribution.name or '')) with tf.name_scope(name) as name: self._distribution = distribution self._bijector = bijector self._kwargs_split_fn = (_default_kwargs_split_fn if kwargs_split_fn is None else kwargs_split_fn) # For convenience we define some handy constants. self._zero = tf.constant(0, dtype=tf.int32, name='zero') self._empty = tf.constant([], dtype=tf.int32, name='empty') # We don't just want to check isinstance(JointDistribution) because # TransformedDistributions with multipart bijectors are effectively # joint but don't inherit from JD. The 'duck-type' test is that # JDs have a structured dtype. self._base_is_joint = tf.nest.is_nested(self.distribution.dtype) if self._base_is_joint: if batch_shape: raise ValueError( 'Overriding the batch shape of a joint distribution' ' ({}) is not currently supported.'.format( self.distribution)) if event_shape: raise ValueError( 'Overriding the event shape of a joint distribution' ' ({}) is not currently supported.'.format( self.distribution)) override_batch_shape = self._empty if batch_shape is None else batch_shape self._override_batch_shape = tensor_util.convert_nonref_to_tensor( override_batch_shape, dtype=tf.int32, name='override_batch_shape') override_event_shape = self._empty if event_shape is None else event_shape self._override_event_shape = tensor_util.convert_nonref_to_tensor( override_event_shape, dtype=tf.int32, name='override_event_shape') # `_is_maybe_{batch, event}_override` is False if we know statically that # the batch/event shape is not being overridden; otherwise it is True. self._is_maybe_event_override = tensorshape_util.dims( self._override_event_shape.shape) != [0] self._is_maybe_batch_override = tensorshape_util.dims( self._override_batch_shape.shape) != [0] dtype = self.bijector.forward_dtype(self.distribution.dtype) self._is_joint = tf.nest.is_nested(dtype) super(TransformedDistribution, self).__init__( dtype=dtype, reparameterization_type=self._distribution. reparameterization_type, validate_args=validate_args, allow_nan_stats=self._distribution.allow_nan_stats, parameters=parameters, name=name)
def __init__(self, kernel, index_points=None, observation_index_points=None, observations=None, observation_noise_variance=0., predictive_noise_variance=None, mean_fn=None, cholesky_fn=None, jitter=1e-6, validate_args=False, allow_nan_stats=False, name='GaussianProcessRegressionModel', _conditional_kernel=None, _conditional_mean_fn=None): """Construct a GaussianProcessRegressionModel instance. Args: kernel: `PositiveSemidefiniteKernel`-like instance representing the GP's covariance function. index_points: `float` `Tensor` representing finite collection, or batch of collections, of points in the index set over which the GP is defined. Shape has the form `[b1, ..., bB, e, f1, ..., fF]` where `F` is the number of feature dimensions and must equal `kernel.feature_ndims` and `e` is the number (size) of index points in each batch. Ultimately this distribution corresponds to an `e`-dimensional multivariate normal. The batch shape must be broadcastable with `kernel.batch_shape` and any batch dims yielded by `mean_fn`. observation_index_points: `float` `Tensor` representing finite collection, or batch of collections, of points in the index set for which some data has been observed. Shape has the form `[b1, ..., bB, e, f1, ..., fF]` where `F` is the number of feature dimensions and must equal `kernel.feature_ndims`, and `e` is the number (size) of index points in each batch. `[b1, ..., bB, e]` must be broadcastable with the shape of `observations`, and `[b1, ..., bB]` must be broadcastable with the shapes of all other batched parameters (`kernel.batch_shape`, `index_points`, etc). The default value is `None`, which corresponds to the empty set of observations, and simply results in the prior predictive model (a GP with noise of variance `predictive_noise_variance`). observations: `float` `Tensor` representing collection, or batch of collections, of observations corresponding to `observation_index_points`. Shape has the form `[b1, ..., bB, e]`, which must be brodcastable with the batch and example shapes of `observation_index_points`. The batch shape `[b1, ..., bB]` must be broadcastable with the shapes of all other batched parameters (`kernel.batch_shape`, `index_points`, etc.). The default value is `None`, which corresponds to the empty set of observations, and simply results in the prior predictive model (a GP with noise of variance `predictive_noise_variance`). observation_noise_variance: `float` `Tensor` representing the variance of the noise in the Normal likelihood distribution of the model. May be batched, in which case the batch shape must be broadcastable with the shapes of all other batched parameters (`kernel.batch_shape`, `index_points`, etc.). Default value: `0.` predictive_noise_variance: `float` `Tensor` representing the variance in the posterior predictive model. If `None`, we simply re-use `observation_noise_variance` for the posterior predictive noise. If set explicitly, however, we use this value. This allows us, for example, to omit predictive noise variance (by setting this to zero) to obtain noiseless posterior predictions of function values, conditioned on noisy observations. mean_fn: Python `callable` that acts on `index_points` to produce a collection, or batch of collections, of mean values at `index_points`. Takes a `Tensor` of shape `[b1, ..., bB, f1, ..., fF]` and returns a `Tensor` whose shape is broadcastable with `[b1, ..., bB]`. Default value: `None` implies the constant zero function. cholesky_fn: Callable which takes a single (batch) matrix argument and returns a Cholesky-like lower triangular factor. Default value: `None`, in which case `make_cholesky_with_jitter_fn` is used with the `jitter` parameter. jitter: `float` scalar `Tensor` added to the diagonal of the covariance matrix to ensure positive definiteness of the covariance matrix. This argument is ignored if `cholesky_fn` is set. Default value: `1e-6`. validate_args: Python `bool`, default `False`. When `True` distribution parameters are checked for validity despite possibly degrading runtime performance. When `False` invalid inputs may silently render incorrect outputs. Default value: `False`. allow_nan_stats: Python `bool`, default `True`. When `True`, statistics (e.g., mean, mode, variance) use the value `NaN` to indicate the result is undefined. When `False`, an exception is raised if one or more of the statistic's batch members are undefined. Default value: `False`. name: Python `str` name prefixed to Ops created by this class. Default value: 'GaussianProcessRegressionModel'. _conditional_kernel: Internal parameter -- do not use. _conditional_mean_fn: Internal parameter -- do not use. Raises: ValueError: if either - only one of `observations` and `observation_index_points` is given, or - `mean_fn` is not `None` and not callable. """ parameters = dict(locals()) with tf.name_scope(name) as name: dtype = dtype_util.common_dtype([ index_points, observation_index_points, observations, observation_noise_variance, predictive_noise_variance, jitter ], tf.float32) index_points = tensor_util.convert_nonref_to_tensor( index_points, dtype=dtype, name='index_points') observation_index_points = tensor_util.convert_nonref_to_tensor( observation_index_points, dtype=dtype, name='observation_index_points') observations = tensor_util.convert_nonref_to_tensor( observations, dtype=dtype, name='observations') observation_noise_variance = tensor_util.convert_nonref_to_tensor( observation_noise_variance, dtype=dtype, name='observation_noise_variance') predictive_noise_variance = tensor_util.convert_nonref_to_tensor( predictive_noise_variance, dtype=dtype, name='predictive_noise_variance') if predictive_noise_variance is None: predictive_noise_variance = observation_noise_variance jitter = tensor_util.convert_nonref_to_tensor(jitter, dtype=dtype, name='jitter') if (observation_index_points is None) != (observations is None): raise ValueError( '`observations` and `observation_index_points` must both be given ' 'or None. Got {} and {}, respectively.'.format( observations, observation_index_points)) # Default to a constant zero function, borrowing the dtype from # index_points to ensure consistency. if mean_fn is None: mean_fn = lambda x: tf.zeros([1], dtype=dtype) else: if not callable(mean_fn): raise ValueError('`mean_fn` must be a Python callable') if cholesky_fn is None: cholesky_fn = cholesky_util.make_cholesky_with_jitter_fn( jitter) self._name = name self._observation_index_points = observation_index_points self._observations = observations self._observation_noise_variance = observation_noise_variance self._predictive_noise_variance = predictive_noise_variance self._jitter = jitter self._validate_args = validate_args with tf.name_scope('init'): if _conditional_kernel is None: _conditional_kernel = tfpk.SchurComplement( base_kernel=kernel, fixed_inputs=observation_index_points, cholesky_fn=cholesky_fn, diag_shift=observation_noise_variance) # Special logic for mean_fn only; SchurComplement already handles the # case of empty observations (ie, falls back to base_kernel). if _is_empty_observation_data( feature_ndims=kernel.feature_ndims, observation_index_points=observation_index_points, observations=observations): if _conditional_mean_fn is None: _conditional_mean_fn = mean_fn else: _validate_observation_data( kernel=kernel, observation_index_points=observation_index_points, observations=observations) if _conditional_mean_fn is None: def conditional_mean_fn(x): """Conditional mean.""" observations = tf.convert_to_tensor( self._observations) observation_index_points = tf.convert_to_tensor( self._observation_index_points) k_x_obs_linop = tf.linalg.LinearOperatorFullMatrix( kernel.matrix(x, observation_index_points)) chol_linop = tf.linalg.LinearOperatorLowerTriangular( _conditional_kernel.divisor_matrix_cholesky( fixed_inputs=observation_index_points)) diff = observations - mean_fn( observation_index_points) return mean_fn(x) + k_x_obs_linop.matvec( chol_linop.solvevec(chol_linop.solvevec(diff), adjoint=True)) _conditional_mean_fn = conditional_mean_fn super(GaussianProcessRegressionModel, self).__init__( kernel=_conditional_kernel, mean_fn=_conditional_mean_fn, index_points=index_points, cholesky_fn=cholesky_fn, jitter=jitter, # What the GP super class calls "observation noise variance" we call # here the "predictive noise variance". We use the observation noise # variance for the fit/solve process above, and predictive for # downstream computations like sampling. observation_noise_variance=predictive_noise_variance, validate_args=validate_args, allow_nan_stats=allow_nan_stats, name=name) self._parameters = parameters
def __init__(self, kernel, index_points=None, mean_fn=None, observation_noise_variance=0., marginal_fn=None, cholesky_fn=None, jitter=1e-6, always_yield_multivariate_normal=False, validate_args=False, allow_nan_stats=False, parameters=None, name='GaussianProcess', _check_marginal_cholesky_fn=True): """Instantiate a GaussianProcess Distribution. Args: kernel: `PositiveSemidefiniteKernel`-like instance representing the GP's covariance function. index_points: `float` `Tensor` representing finite (batch of) vector(s) of points in the index set over which the GP is defined. Shape has the form `[b1, ..., bB, e, f1, ..., fF]` where `F` is the number of feature dimensions and must equal `kernel.feature_ndims` and `e` is the number (size) of index points in each batch. Ultimately this distribution corresponds to a `e`-dimensional multivariate normal. The batch shape must be broadcastable with `kernel.batch_shape` and any batch dims yielded by `mean_fn`. mean_fn: Python `callable` that acts on `index_points` to produce a (batch of) vector(s) of mean values at `index_points`. Takes a `Tensor` of shape `[b1, ..., bB, f1, ..., fF]` and returns a `Tensor` whose shape is broadcastable with `[b1, ..., bB]`. Default value: `None` implies constant zero function. observation_noise_variance: `float` `Tensor` representing (batch of) scalar variance(s) of the noise in the Normal likelihood distribution of the model. If batched, the batch shape must be broadcastable with the shapes of all other batched parameters (`kernel.batch_shape`, `index_points`, etc.). Default value: `0.` marginal_fn: A Python callable that takes a location, covariance matrix, optional `validate_args`, `allow_nan_stats` and `name` arguments, and returns a multivariate normal subclass of `tfd.Distribution`. At most one of `cholesky_fn` and `marginal_fn` should be set. Default value: `None`, in which case a Cholesky-factorizing function is created using `make_cholesky_factored_marginal_fn` and the `cholesky_fn` argument. cholesky_fn: Callable which takes a single (batch) matrix argument and returns a Cholesky-like lower triangular factor. Default value: `None`, in which case `make_cholesky_with_jitter_fn` is used with the `jitter` parameter. At most one of `cholesky_fn` and `marginal_fn` should be set. jitter: `float` scalar `Tensor` added to the diagonal of the covariance matrix to ensure positive definiteness of the covariance matrix, when `marginal_fn` and `cholesky_fn` is None. This argument is ignored if `cholesky_fn` is set. Default value: `1e-6`. always_yield_multivariate_normal: If `False` (the default), we produce a scalar `Normal` distribution when the number of `index_points` is statically known to be `1`. If `True`, we avoid this behavior, ensuring that the event shape will retain the `1` from `index_points`. validate_args: Python `bool`, default `False`. When `True` distribution parameters are checked for validity despite possibly degrading runtime performance. When `False` invalid inputs may silently render incorrect outputs. Default value: `False`. allow_nan_stats: Python `bool`, default `True`. When `True`, statistics (e.g., mean, mode, variance) use the value "`NaN`" to indicate the result is undefined. When `False`, an exception is raised if one or more of the statistic's batch members are undefined. Default value: `False`. parameters: For subclasses, a dict of constructor arguments. name: Python `str` name prefixed to Ops created by this class. Default value: "GaussianProcess". _check_marginal_cholesky_fn: Internal parameter -- do not use. Raises: ValueError: if `mean_fn` is not `None` and is not callable. """ parameters = dict(locals()) if parameters is None else parameters with tf.name_scope(name) as name: dtype = dtype_util.common_dtype( { 'index_points': index_points, 'observation_noise_variance': observation_noise_variance, 'jitter': jitter }, tf.float32) index_points = tensor_util.convert_nonref_to_tensor( index_points, dtype=dtype, name='index_points') jitter = tensor_util.convert_nonref_to_tensor(jitter, dtype=dtype, name='jitter') observation_noise_variance = tensor_util.convert_nonref_to_tensor( observation_noise_variance, dtype=dtype, name='observation_noise_variance') self._kernel = kernel self._index_points = index_points # Default to a constant zero function, borrowing the dtype from # index_points to ensure consistency. if mean_fn is None: mean_fn = lambda x: tf.zeros([1], dtype=dtype) else: if not callable(mean_fn): raise ValueError('`mean_fn` must be a Python callable') self._mean_fn = mean_fn self._observation_noise_variance = observation_noise_variance self._jitter = jitter self._cholesky_fn = cholesky_fn if (_check_marginal_cholesky_fn and marginal_fn is not None and cholesky_fn is not None): raise ValueError( 'At most one of `marginal_fn` and `cholesky_fn` should be set.' ) if marginal_fn is None: if cholesky_fn is None: self._cholesky_fn = cholesky_util.make_cholesky_with_jitter_fn( jitter) self._marginal_fn = make_cholesky_factored_marginal_fn( self._cholesky_fn) else: self._marginal_fn = marginal_fn self._always_yield_multivariate_normal = always_yield_multivariate_normal with tf.name_scope('init'): super(GaussianProcess, self).__init__( dtype=dtype, reparameterization_type=reparameterization. FULLY_REPARAMETERIZED, validate_args=validate_args, allow_nan_stats=allow_nan_stats, parameters=parameters, name=name)
def __init__(self, bin_widths, bin_heights, knot_slopes, range_min=-1, validate_args=False, name=None): """Construct a new RationalQuadraticSpline bijector. For each argument, the innermost axis indexes bins/knots and batch axes index axes of `x`/`y` spaces. A `RationalQuadraticSpline` with a separate transform for each of three dimensions might have `bin_widths` shaped `[3, 32]`. To use the same spline for each of `x`'s three dimensions we may broadcast against `x` and use a `bin_widths` parameter shaped `[32]`. Parameters will be broadcast against each other and against the input `x`/`y`s, so if we want fixed slopes, we can use kwarg `knot_slopes=1`. A typical recipe for acquiring compatible bin widths and heights would be: ```python nbins = unconstrained_vector.shape[-1] range_min, range_max, min_bin_size = -1, 1, 1e-2 scale = range_max - range_min - nbins * min_bin_size bin_widths = tf.math.softmax(unconstrained_vector) * scale + min_bin_size ``` Args: bin_widths: The widths of the spans between subsequent knot `x` positions, a floating point `Tensor`. Must be positive, and at least 1-D. Innermost axis must sum to the same value as `bin_heights`. The knot `x` positions will be a first at `range_min`, followed by knots at `range_min + cumsum(bin_widths, axis=-1)`. bin_heights: The heights of the spans between subsequent knot `y` positions, a floating point `Tensor`. Must be positive, and at least 1-D. Innermost axis must sum to the same value as `bin_widths`. The knot `y` positions will be a first at `range_min`, followed by knots at `range_min + cumsum(bin_heights, axis=-1)`. knot_slopes: The slope of the spline at each knot, a floating point `Tensor`. Must be positive. `1`s are implicitly padded for the first and last implicit knots corresponding to `range_min` and `range_min + sum(bin_widths, axis=-1)`. Innermost axis size should be 1 less than that of `bin_widths`/`bin_heights`, or 1 for broadcasting. range_min: The `x`/`y` position of the first knot, which has implicit slope `1`. `range_max` is implicit, and can be computed as `range_min + sum(bin_widths, axis=-1)`. Scalar floating point `Tensor`. validate_args: Toggles argument validation (can hurt performance). name: Optional name scope for associated ops. (Defaults to `'RationalQuadraticSpline'`). """ parameters = dict(locals()) with tf.name_scope(name or 'RationalQuadraticSpline') as name: dtype = dtype_util.common_dtype( [bin_widths, bin_heights, knot_slopes, range_min], dtype_hint=tf.float32) self._bin_widths = tensor_util.convert_nonref_to_tensor( bin_widths, dtype=dtype, name='bin_widths') self._bin_heights = tensor_util.convert_nonref_to_tensor( bin_heights, dtype=dtype, name='bin_heights') self._knot_slopes = tensor_util.convert_nonref_to_tensor( knot_slopes, dtype=dtype, name='knot_slopes') self._range_min = tensor_util.convert_nonref_to_tensor( range_min, dtype=dtype, name='range_min') super(RationalQuadraticSpline, self).__init__(dtype=dtype, forward_min_event_ndims=0, validate_args=validate_args, parameters=parameters, name=name)
def __init__(self, loc, scale, quadrature_size=8, quadrature_fn=quadrature_scheme_lognormal_quantiles, validate_args=False, allow_nan_stats=True, name='PoissonLogNormalQuadratureCompound'): """Constructs the PoissonLogNormalQuadratureCompound`. Note: `probs` returned by (optional) `quadrature_fn` are presumed to be either a length-`quadrature_size` vector or a batch of vectors in 1-to-1 correspondence with the returned `grid`. (I.e., broadcasting is only partially supported.) Args: loc: `float`-like (batch of) scalar `Tensor`; the location parameter of the LogNormal prior. scale: `float`-like (batch of) scalar `Tensor`; the scale parameter of the LogNormal prior. quadrature_size: Python `int` scalar representing the number of quadrature points. quadrature_fn: Python callable taking `loc`, `scale`, `quadrature_size`, `validate_args` and returning `tuple(grid, probs)` representing the LogNormal grid and corresponding normalized weight. Default value: `quadrature_scheme_lognormal_quantiles`. validate_args: Python `bool`, default `False`. When `True` distribution parameters are checked for validity despite possibly degrading runtime performance. When `False` invalid inputs may silently render incorrect outputs. allow_nan_stats: Python `bool`, default `True`. When `True`, statistics (e.g., mean, mode, variance) use the value '`NaN`' to indicate the result is undefined. When `False`, an exception is raised if one or more of the statistic's batch members are undefined. name: Python `str` name prefixed to Ops created by this class. Raises: TypeError: if `quadrature_grid` and `quadrature_probs` have different base `dtype`. """ parameters = dict(locals()) with tf.name_scope(name) as name: dtype = dtype_util.common_dtype([loc, scale], tf.float32) self._loc = tensor_util.convert_nonref_to_tensor(loc, name='loc', dtype=dtype) self._scale = tensor_util.convert_nonref_to_tensor(scale, name='scale', dtype=dtype) self._quadrature_fn = quadrature_fn dtype_util.assert_same_float_dtype([self._loc, self._scale]) self._quadrature_size = quadrature_size super(PoissonLogNormalQuadratureCompound, self).__init__( dtype=dtype, reparameterization_type=reparameterization.NOT_REPARAMETERIZED, validate_args=validate_args, allow_nan_stats=allow_nan_stats, parameters=parameters, name=name)
def __init__(self, paddings=((0, 1), ), mode='CONSTANT', constant_values=0, axis=None, validate_args=False, name=None): """Initializes the `Pad` bijector. Args: paddings: A vector-shaped `Tensor` of `int` pairs representing the number of elements to pad on the left and right, respectively. Default value: `((0, 1),)`. mode: One of `'CONSTANT'`, `'REFLECT'`, or `'SYMMETRIC'` (case-insensitive). For more details, see `tf.pad`. constant_values: In "CONSTANT" mode, the scalar pad value to use. Must be same type as `tensor`. For more details, see `tf.pad`. axis: The dimensions for which `paddings` are applied. Must be 1:1 with `paddings` or `None`. Default value: `None` (i.e., `tf.range(start=-len(paddings), limit=0)`). validate_args: Python `bool` indicating whether arguments should be checked for correctness. Default value: `False`. name: Python `str`, name given to ops managed by this object. Default value: `None` (i.e., `'pad'`). """ parameters = dict(locals()) with tf.name_scope(name or 'pad') as name: paddings = tensor_util.convert_nonref_to_tensor( paddings, dtype_hint=tf.int32, name='paddings', as_shape_tensor=True) if axis is None: axis = ps.range(start=-ps.size0(paddings), limit=0, dtype=tf.int32, name='axis') else: axis = tensor_util.convert_nonref_to_tensor( axis, dtype_hint=tf.int32, name='axis', as_shape_tensor=True) axis_ = tf.get_static_value(axis) if axis_ is None: raise NotImplementedError( 'Argument `axis` must be known statically. If you need this ' 'feature, please contact `[email protected]`.') self._axis = axis self._paddings = paddings self._mode = mode self._constant_values = tensor_util.convert_nonref_to_tensor( constant_values, dtype_hint=tf.float32, name='constant_values') min_event_ndims_ = int(-np.min( np.pad(np.reshape(axis_, newshape=[-1]), mode='constant', pad_width=[[0, 1]]))) super(Pad, self).__init__(forward_min_event_ndims=min_event_ndims_, inverse_min_event_ndims=min_event_ndims_, is_constant_jacobian=True, validate_args=validate_args, parameters=parameters, name=name)
def __init__(self, loc=None, scale_tril=None, validate_args=False, allow_nan_stats=True, name='MultivariateNormalTriL'): """Construct Multivariate Normal distribution on `R^k`. The `batch_shape` is the broadcast shape between `loc` and `scale` arguments. The `event_shape` is given by last dimension of the matrix implied by `scale`. The last dimension of `loc` (if provided) must broadcast with this. Recall that `covariance = scale @ scale.T`. A (non-batch) `scale` matrix is: ```none scale = scale_tril ``` where `scale_tril` is lower-triangular `k x k` matrix with non-zero diagonal, i.e., `tf.diag_part(scale_tril) != 0`. Additional leading dimensions (if any) will index batches. Args: loc: Floating-point `Tensor`. If this is set to `None`, `loc` is implicitly `0`. When specified, may have shape `[B1, ..., Bb, k]` where `b >= 0` and `k` is the event size. scale_tril: Floating-point, lower-triangular `Tensor` with non-zero diagonal elements. `scale_tril` has shape `[B1, ..., Bb, k, k]` where `b >= 0` and `k` is the event size. validate_args: Python `bool`, default `False`. When `True` distribution parameters are checked for validity despite possibly degrading runtime performance. When `False` invalid inputs may silently render incorrect outputs. allow_nan_stats: Python `bool`, default `True`. When `True`, statistics (e.g., mean, mode, variance) use the value "`NaN`" to indicate the result is undefined. When `False`, an exception is raised if one or more of the statistic's batch members are undefined. name: Python `str` name prefixed to Ops created by this class. Raises: ValueError: if neither `loc` nor `scale_tril` are specified. """ parameters = dict(locals()) if loc is None and scale_tril is None: raise ValueError( 'Must specify one or both of `loc`, `scale_tril`.') with tf.name_scope(name) as name: dtype = dtype_util.common_dtype([loc, scale_tril], tf.float32) loc = tensor_util.convert_nonref_to_tensor(loc, name='loc', dtype=dtype) scale_tril = tensor_util.convert_nonref_to_tensor( scale_tril, name='scale_tril', dtype=dtype) self._scale_tril = scale_tril if scale_tril is None: scale = tf.linalg.LinearOperatorIdentity( num_rows=distribution_util.dimension_size(loc, -1), dtype=loc.dtype, is_self_adjoint=True, is_positive_definite=True, assert_proper_shapes=validate_args) else: # No need to validate that scale_tril is non-singular. # LinearOperatorLowerTriangular has an assert_non_singular # method that is called by the Bijector. scale = tf.linalg.LinearOperatorLowerTriangular( scale_tril, is_non_singular=True, is_self_adjoint=False, is_positive_definite=False) super(MultivariateNormalTriL, self).__init__(loc=loc, scale=scale, validate_args=validate_args, allow_nan_stats=allow_nan_stats, name=name) self._parameters = parameters
def __init__(self, outcomes, logits=None, probs=None, rtol=None, atol=None, validate_args=False, allow_nan_stats=True, name='FiniteDiscrete'): """Construct a finite discrete contribution. Args: outcomes: A 1-D floating or integer `Tensor`, representing a list of possible outcomes in strictly ascending order. logits: A floating N-D `Tensor`, `N >= 1`, representing the log probabilities of a set of FiniteDiscrete distributions. The first `N - 1` dimensions index into a batch of independent distributions and the last dimension represents a vector of logits for each discrete value. Only one of `logits` or `probs` should be passed in. probs: A floating N-D `Tensor`, `N >= 1`, representing the probabilities of a set of FiniteDiscrete distributions. The first `N - 1` dimensions index into a batch of independent distributions and the last dimension represents a vector of probabilities for each discrete value. Only one of `logits` or `probs` should be passed in. rtol: `Tensor` with same `dtype` as `outcomes`. The relative tolerance for floating number comparison. Only effective when `outcomes` is a floating `Tensor`. Default is `10 * eps`. atol: `Tensor` with same `dtype` as `outcomes`. The absolute tolerance for floating number comparison. Only effective when `outcomes` is a floating `Tensor`. Default is `10 * eps`. validate_args: Python `bool`, default `False`. When `True` distribution parameters are checked for validity despite possibly degrading runtime performance. When `False` invalid inputs may render incorrect outputs. allow_nan_stats: Python `bool`, default `True`. When `True`, statistics (e.g., mean, mode, variance) use the value '`NaN`' to indicate the result is undefined. When `False`, an exception is raised if one or more of the statistic's batch members are undefined. name: Python `str` name prefixed to Ops created by this class. """ parameters = dict(locals()) with tf.name_scope(name) as name: outcomes_dtype = dtype_util.common_dtype([outcomes], dtype_hint=tf.float32) self._outcomes = tensor_util.convert_nonref_to_tensor( outcomes, dtype_hint=outcomes_dtype, name='outcomes') if dtype_util.is_floating(self._outcomes.dtype): eps = np.finfo(dtype_util.as_numpy_dtype(outcomes_dtype)).eps self._rtol = 10 * eps if rtol is None else rtol self._atol = 10 * eps if atol is None else atol else: self._rtol = None self._atol = None self._categorical = categorical.Categorical( logits=logits, probs=probs, dtype=tf.int32, validate_args=validate_args, allow_nan_stats=allow_nan_stats) super(FiniteDiscrete, self).__init__( dtype=self._outcomes.dtype, reparameterization_type=reparameterization.NOT_REPARAMETERIZED, validate_args=validate_args, allow_nan_stats=allow_nan_stats, parameters=parameters, name=name)
def __init__(self, event_shape_out, event_shape_in=(-1, ), validate_args=False, name=None): """Creates a `Reshape` bijector. Args: event_shape_out: An `int`-like vector-shaped `Tensor` representing the event shape of the transformed output. event_shape_in: An optional `int`-like vector-shape `Tensor` representing the event shape of the input. This is required in order to define inverse operations; the default of (-1,) assumes a vector-shaped input. validate_args: Python `bool` indicating whether arguments should be checked for correctness. name: Python `str`, name given to ops managed by this object. Raises: TypeError: if either `event_shape_in` or `event_shape_out` has non-integer `dtype`. ValueError: if either of `event_shape_in` or `event_shape_out` has non-vector shape (`rank > 1`), or if their sizes do not match. """ parameters = dict(locals()) with tf.name_scope(name or 'reshape') as name: dtype = dtype_util.common_dtype([event_shape_out, event_shape_in], dtype_hint=tf.int32) event_shape_out = tensor_util.convert_nonref_to_tensor( event_shape_out, name='event_shape_out', dtype=dtype, as_shape_tensor=True) event_shape_in = tensor_util.convert_nonref_to_tensor( event_shape_in, name='event_shape_in', dtype=dtype, as_shape_tensor=True) forward_min_event_ndims_ = _rank_from_shape(event_shape_in) if forward_min_event_ndims_ is None: raise NotImplementedError( 'The length of `event_shape_in` must be statically known. For ' 'dynamic support, please contact `[email protected]`.' ) inverse_min_event_ndims_ = _rank_from_shape(event_shape_out) if inverse_min_event_ndims_ is None: raise NotImplementedError( 'The length of `event_shape_out` must be statically known. For ' 'dynamic support, please contact `[email protected]`.' ) self._event_shape_in = event_shape_in self._event_shape_out = event_shape_out super(Reshape, self).__init__( forward_min_event_ndims=forward_min_event_ndims_, inverse_min_event_ndims=inverse_min_event_ndims_, is_constant_jacobian=True, validate_args=validate_args, parameters=parameters, name=name or 'reshape')
def __init__(self, perm=None, rightmost_transposed_ndims=None, validate_args=False, name='transpose'): """Instantiates the `Transpose` bijector. Args: perm: Positive `int32` vector-shaped `Tensor` representing permutation of rightmost dims (for forward transformation). Note that the `0`th index represents the first of the rightmost dims and the largest value must be `rightmost_transposed_ndims - 1` and corresponds to `tf.rank(x) - 1`. Only one of `perm` and `rightmost_transposed_ndims` can (and must) be specified. The number of elements in a permutation must have a value that can be determined statically. Default value: `tf.range(start=rightmost_transposed_ndims, limit=-1, delta=-1)`. rightmost_transposed_ndims: Positive `int32` scalar-shaped `Tensor` representing the number of rightmost dimensions to permute. Only one of `perm` and `rightmost_transposed_ndims` can (and must) be specified. If `rightmost_transposed_ndims` is specified, the rightmost dims are reversed. This argument must have a value that can be determined statically. Default value: `tf.size(perm)`. validate_args: Python `bool` indicating whether arguments should be checked for correctness. name: Python `str` name given to ops managed by this object. Raises: ValueError: if both or neither `perm` and `rightmost_transposed_ndims` are specified. NotImplementedError: if `rightmost_transposed_ndims` is not known prior to graph execution. """ with tf.name_scope(name) as name: # We need to determine `forward_min_event_ndims` statically, which # requires that we know `rightmost_transposed_ndims` statically. # So the corresponding assertions go here rather than in # `_parameter_control_dependencies` if (rightmost_transposed_ndims is None) == (perm is None): raise ValueError('Must specify exactly one of ' '`rightmost_transposed_ndims` and `perm`.') if rightmost_transposed_ndims is not None: rightmost_transposed_ndims = tensor_util.convert_nonref_to_tensor( rightmost_transposed_ndims, dtype_hint=np.int32) if not dtype_util.is_integer(rightmost_transposed_ndims.dtype): raise TypeError( '`rightmost_transposed_ndims` must be integer type.') rightmost_transposed_ndims_ = tf.get_static_value( rightmost_transposed_ndims) if rightmost_transposed_ndims_ is None: raise NotImplementedError( '`rightmost_transposed_ndims` must be ' 'known prior to graph execution.') msg = '`rightmost_transposed_ndims` must be non-negative.' if rightmost_transposed_ndims_ < 0: raise ValueError( msg[:-1] + ', saw: {}.'.format(rightmost_transposed_ndims_)) perm_start = (distribution_util.prefer_static_value( rightmost_transposed_ndims) - 1) perm = tf.range(start=perm_start, limit=-1, delta=-1, name='perm') else: # perm is not None: perm = tensor_util.convert_nonref_to_tensor( perm, dtype_hint=np.int32, name='perm') rightmost_transposed_ndims = tf.size( perm, name='rightmost_transposed_ndims') rightmost_transposed_ndims_ = tf.get_static_value( rightmost_transposed_ndims) # TODO(b/110828604): If bijector base class ever supports dynamic # `min_event_ndims`, then this class already works dynamically and the # following five lines can be removed. if rightmost_transposed_ndims_ is None: raise NotImplementedError( '`rightmost_transposed_ndims` must be ' 'known prior to graph execution.') else: rightmost_transposed_ndims_ = int(rightmost_transposed_ndims_) self._perm = perm self._rightmost_transposed_ndims = rightmost_transposed_ndims self._initial_rightmost_transposed_ndims = rightmost_transposed_ndims_ super(Transpose, self).__init__( forward_min_event_ndims=rightmost_transposed_ndims_, is_constant_jacobian=True, validate_args=validate_args, name=name)
def __init__(self, kernel, x1, x2=None, num_matmul_parts=None, is_non_singular=None, is_self_adjoint=None, is_positive_definite=None, is_square=None, name=None): """Initializes the LinearOperator. This object implicitly represents the covariance matrix of `x1` and `x2` (`x1` if `x2` not provided) under the given `kernel`. This operator assumes one example dimension on each set of index points, which indexes the corresponding axis of the matrix. All outer dimensions are considered batch dimensions. Use this to avoid materializing the full matrix for such operations as: - accessing diagonal (`diag_part` method) - accessing a [batch of] row[s] (`row` method) - accessing a [batch of] column[s] (`col` method) Use this to perform matrix-vector products on very large covariance matrices by chunking the covariance matrix into parts, computing just that part of the output, computing a part of the covariance matrix, and the matrix-vector product, then forgetting the partial covariance matrix. Internally, uses recomputed gradients to avoid retaining infeasibly-large intermediate activations. Args: kernel: A `tfp.math.psd_kernels.PositiveSemidefiniteKernel` instance. x1: A floating point `Tensor`, the row index points. x2: Optional, a floating point `Tensor`, the column index points. If not provided, uses `x1`. num_matmul_parts: An optional Python `int` specifying the number of partitions into which the matrix should be broken when applying this linear operator. (An extra remainder partition is implied for uneven partitioning.) Because the use-case is avoiding a memory blowup, the partitioned matmul uses `parallel_iterations=1` and `back_prop=False`. is_non_singular: Expect that this operator is non-singular. is_self_adjoint: Expect that this operator is equal to its hermitian transpose. If `dtype` is real, this is equivalent to being symmetric. is_positive_definite: Expect that this operator is positive definite, meaning the quadratic form `x^H A x` has positive real part for all nonzero `x`. Note that we do not require the operator to be self-adjoint to be positive-definite. See: https://en.wikipedia.org/wiki/Positive-definite_matrix#Extension_for_non-symmetric_matrices is_square: Expect that this operator acts like square [batch] matrices. name: Optional name for related ops. """ dtype = dtype_util.common_dtype([kernel, x1, x2], dtype_hint=tf.float64) self._kernel = kernel self._x1 = tensor_util.convert_nonref_to_tensor(x1, dtype=dtype) self._x2 = tensor_util.convert_nonref_to_tensor(x2, dtype=dtype) self._num_matmul_parts = tensor_util.convert_nonref_to_tensor( num_matmul_parts, dtype=tf.int32) if self._x2 is None: if is_non_singular is not None and not is_non_singular: raise ValueError( 'Operator is non-singular with a single set of index points.' ) is_non_singular = True if is_self_adjoint is not None and not is_self_adjoint: raise ValueError( 'Operator is self-adjoint with a single set of index points.' ) is_self_adjoint = True if is_positive_definite is not None and not is_positive_definite: raise ValueError( 'Operator is positive-definite with a single set of index points.' ) is_positive_definite = True if is_square is not None and not is_square: raise ValueError( 'Operator is square with a single set of index points.') is_square = True super(LinearOperatorPSDKernel, self).__init__(dtype, is_non_singular=is_non_singular, is_self_adjoint=is_self_adjoint, is_positive_definite=is_positive_definite, is_square=is_square, name=name or 'LinearOperatorPSDKernel')
def __init__(self, x, name=None): super(FakeModule, self).__init__(name) self.x = tensor_util.convert_nonref_to_tensor(x)
def __init__(self, loc, scale, validate_args=False, allow_nan_stats=True, name='Gumbel'): """Construct Gumbel distributions with location and scale `loc` and `scale`. The parameters `loc` and `scale` must be shaped in a way that supports broadcasting (e.g. `loc + scale` is a valid operation). Args: loc: Floating point tensor, the means of the distribution(s). scale: Floating point tensor, the scales of the distribution(s). scale must contain only positive values. validate_args: Python `bool`, default `False`. When `True` distribution parameters are checked for validity despite possibly degrading runtime performance. When `False` invalid inputs may silently render incorrect outputs. Default value: `False`. allow_nan_stats: Python `bool`, default `True`. When `True`, statistics (e.g., mean, mode, variance) use the value `NaN` to indicate the result is undefined. When `False`, an exception is raised if one or more of the statistic's batch members are undefined. Default value: `True`. name: Python `str` name prefixed to Ops created by this class. Default value: `'Gumbel'`. Raises: TypeError: if loc and scale are different dtypes. """ parameters = dict(locals()) with tf.name_scope(name) as name: dtype = dtype_util.common_dtype([loc, scale], dtype_hint=tf.float32) loc = tensor_util.convert_nonref_to_tensor(loc, name='loc', dtype=dtype) scale = tensor_util.convert_nonref_to_tensor(scale, name='scale', dtype=dtype) dtype_util.assert_same_float_dtype([loc, scale]) # Positive scale is asserted by the incorporated Gumbel bijector. self._gumbel_bijector = gumbel_cdf_bijector.GumbelCDF( loc=loc, scale=scale, validate_args=validate_args) # Because the uniform sampler generates samples in `[0, 1)` this would # cause samples to lie in `(inf, -inf]` instead of `(inf, -inf)`. To fix # this, we use `np.finfo(dtype_util.as_numpy_dtype(self.dtype).tiny` # because it is the smallest, positive, 'normal' number. super(Gumbel, self).__init__( distribution=uniform.Uniform(low=np.finfo( dtype_util.as_numpy_dtype(dtype)).tiny, high=tf.ones([], dtype=dtype), allow_nan_stats=allow_nan_stats), # The Gumbel bijector encodes the CDF function as the forward, # and hence needs to be inverted. bijector=invert_bijector.Invert(self._gumbel_bijector, validate_args=validate_args), parameters=parameters, name=name)
def test_tf_tensor(self): x = tf.constant(1.) y = tensor_util.convert_nonref_to_tensor(x) self.assertIs(x, y)
def __init__(self, df, scale=None, scale_tril=None, input_output_cholesky=False, validate_args=False, allow_nan_stats=True, name='Wishart'): """Construct Wishart distributions. Args: df: `float` or `double` `Tensor`. Degrees of freedom, must be greater than or equal to dimension of the scale matrix. scale: `float` or `double` `Tensor`. The symmetric positive definite scale matrix of the distribution. Exactly one of `scale` and 'scale_tril` must be passed. scale_tril: `float` or `double` `Tensor`. The Cholesky factorization of the symmetric positive definite scale matrix of the distribution. Exactly one of `scale` and 'scale_tril` must be passed. input_output_cholesky: Python `bool`. If `True`, functions whose input or output have the semantics of samples assume inputs are in Cholesky form and return outputs in Cholesky form. In particular, if this flag is `True`, input to `log_prob` is presumed of Cholesky form and output from `sample`, `mean`, and `mode` are of Cholesky form. Setting this argument to `True` is purely a computational optimization and does not change the underlying distribution; for instance, `mean` returns the Cholesky of the mean, not the mean of Cholesky factors. The `variance` and `stddev` methods are unaffected by this flag. Default value: `False` (i.e., input/output does not have Cholesky semantics). validate_args: Python `bool`, default `False`. When `True` distribution parameters are checked for validity despite possibly degrading runtime performance. When `False` invalid inputs may silently render incorrect outputs. allow_nan_stats: Python `bool`, default `True`. When `True`, statistics (e.g., mean, mode, variance) use the value '`NaN`' to indicate the result is undefined. When `False`, an exception is raised if one or more of the statistic's batch members are undefined. name: Python `str` name prefixed to Ops created by this class. Raises: ValueError: if zero or both of 'scale' and 'scale_tril' are passed in. """ parameters = dict(locals()) with tf.name_scope(name) as name: if (scale is None) == (scale_tril is None): raise ValueError( 'Must pass scale or scale_tril, but not both.') dtype = dtype_util.common_dtype([df, scale, scale_tril], tf.float32) df = tensor_util.convert_nonref_to_tensor(df, name='df', dtype=dtype) if scale is not None: scale = tensor_util.convert_nonref_to_tensor(scale, name='scale', dtype=dtype) self._scale_tril = tf.linalg.cholesky(scale) else: # scale_tril is not None self._scale_tril = tensor_util.convert_nonref_to_tensor( scale_tril, name='scale_tril', dtype=dtype) self._scale_full = scale super(Wishart, self).__init__(df=df, scale=tf.linalg.LinearOperatorLowerTriangular( tril=self._scale_tril, is_non_singular=True, is_positive_definite=True, is_square=True), input_output_cholesky=input_output_cholesky, validate_args=validate_args, allow_nan_stats=allow_nan_stats, name=name) self._parameters = parameters
def __init__(self, power, dtype=tf.int32, interpolate_nondiscrete=True, sample_maximum_iterations=100, validate_args=False, allow_nan_stats=False, name='Zipf'): """Initialize a batch of Zipf distributions. Args: power: `Float` like `Tensor` representing the power parameter. Must be strictly greater than `1`. dtype: The `dtype` of `Tensor` returned by `sample`. Default value: `tf.int32`. interpolate_nondiscrete: Python `bool`. When `False`, `log_prob` returns `-inf` (and `prob` returns `0`) for non-integer inputs. When `True`, `log_prob` evaluates the continuous function `-power log(k) - log(zeta(power))` , which matches the Zipf pmf at integer arguments `k` (note that this function is not itself a normalized probability log-density). Default value: `True`. sample_maximum_iterations: Maximum number of iterations of allowable iterations in `sample`. When `validate_args=True`, samples which fail to reach convergence (subject to this cap) are masked out with `self.dtype.min` or `nan` depending on `self.dtype.is_integer`. Default value: `100`. validate_args: Python `bool`, default `False`. When `True` distribution parameters are checked for validity despite possibly degrading runtime performance. When `False` invalid inputs may silently render incorrect outputs. Default value: `False`. allow_nan_stats: Python `bool`, default `True`. When `True`, statistics (e.g., mean, mode, variance) use the value "`NaN`" to indicate the result is undefined. When `False`, an exception is raised if one or more of the statistic's batch members are undefined. Default value: `False`. name: Python `str` name prefixed to Ops created by this class. Default value: `'Zipf'`. Raises: TypeError: if `power` is not `float` like. """ parameters = dict(locals()) with tf.name_scope(name) as name: self._power = tensor_util.convert_nonref_to_tensor( power, name='power', dtype=dtype_util.common_dtype([power], dtype_hint=tf.float32)) if (not dtype_util.is_floating(self._power.dtype) or dtype_util.base_equal(self._power.dtype, tf.float16)): raise TypeError( 'power.dtype ({}) is not a supported `float` type.'.format( dtype_util.name(self._power.dtype))) self._interpolate_nondiscrete = interpolate_nondiscrete self._sample_maximum_iterations = sample_maximum_iterations super(Zipf, self).__init__( dtype=dtype, reparameterization_type=reparameterization.NOT_REPARAMETERIZED, validate_args=validate_args, allow_nan_stats=allow_nan_stats, parameters=parameters, name=name)
def __init__(self, num_or_size_splits, axis=-1, validate_args=False, name='split'): """Creates the bijector. Args: num_or_size_splits: Either a Python integer indicating the number of splits along `axis` or a 1-D integer `Tensor` or Python list containing the sizes of each output tensor along `axis`. If a list/`Tensor`, it may contain at most one value of `-1`, which indicates a split size that is unknown and determined from input. axis: A negative integer or scalar `int32` `Tensor`. The dimension along which to split. Must be negative to enable the bijector to support arbitrary batch dimensions. Defaults to -1 (note that this is different from the `tf.Split` default of `0`). Must be statically known. validate_args: Python `bool` indicating whether arguments should be checked for correctness. name: Python `str`, name given to ops managed by this object. """ parameters = dict(locals()) with tf.name_scope(name) as name: if isinstance(num_or_size_splits, numbers.Integral): self._num_splits = num_or_size_splits self._split_sizes = None else: self._split_sizes = tensor_util.convert_nonref_to_tensor( num_or_size_splits, name='num_or_size_splits', dtype=tf.int32) if tensorshape_util.rank(self._split_sizes.shape) != 1: raise ValueError( '`num_or_size_splits` must be an integer or 1-D `Tensor`.' ) num_splits = tensorshape_util.as_list( self._split_sizes.shape)[0] if num_splits is None: raise ValueError( 'If `num_or_size_splits` is a vector of split sizes ' 'it must have a statically-known number of ' 'elements.') self._num_splits = num_splits static_axis = tf.get_static_value(axis) if static_axis is None: raise ValueError('`axis` must be statically known.') if static_axis >= 0: raise ValueError( '`axis` must be negative. Got {}'.format(axis)) self._axis = tf.convert_to_tensor(axis, tf.int32) super(Split, self).__init__(forward_min_event_ndims=-axis, inverse_min_event_ndims=[-axis] * self.num_splits, is_constant_jacobian=True, validate_args=validate_args, parameters=parameters, name=name)
def __init__(self, x, y, name=None): with tf.name_scope(name or 'Adder') as name: self._x = tensor_util.convert_nonref_to_tensor(x) self._y = tensor_util.convert_nonref_to_tensor(y) self._name = name
def __init__(self, rate=None, log_rate=None, interpolate_nondiscrete=True, validate_args=False, allow_nan_stats=True, name='Poisson'): """Initialize a batch of Poisson distributions. Args: rate: Floating point tensor, the rate parameter. `rate` must be positive. Must specify exactly one of `rate` and `log_rate`. log_rate: Floating point tensor, the log of the rate parameter. Must specify exactly one of `rate` and `log_rate`. interpolate_nondiscrete: Python `bool`. When `False`, `log_prob` returns `-inf` (and `prob` returns `0`) for non-integer inputs. When `True`, `log_prob` evaluates the continuous function `k * log_rate - lgamma(k+1) - rate`, which matches the Poisson pmf at integer arguments `k` (note that this function is not itself a normalized probability log-density). Default value: `True`. validate_args: Python `bool`. When `True` distribution parameters are checked for validity despite possibly degrading runtime performance. When `False` invalid inputs may silently render incorrect outputs. Default value: `False`. allow_nan_stats: Python `bool`. When `True`, statistics (e.g., mean, mode, variance) use the value "`NaN`" to indicate the result is undefined. When `False`, an exception is raised if one or more of the statistic's batch members are undefined. Default value: `True`. name: Python `str` name prefixed to Ops created by this class. Raises: ValueError: if none or both of `rate`, `log_rate` are specified. TypeError: if `rate` is not a float-type. TypeError: if `log_rate` is not a float-type. """ parameters = dict(locals()) if (rate is None) == (log_rate is None): raise ValueError( 'Must specify exactly one of `rate` and `log_rate`.') with tf.name_scope(name) as name: dtype = dtype_util.common_dtype([rate, log_rate], dtype_hint=tf.float32) if not dtype_util.is_floating(dtype): raise TypeError( '[log_]rate.dtype ({}) is a not a float-type.'.format( dtype_util.name(dtype))) self._rate = tensor_util.convert_nonref_to_tensor(rate, name='rate', dtype=dtype) self._log_rate = tensor_util.convert_nonref_to_tensor( log_rate, name='log_rate', dtype=dtype) self._interpolate_nondiscrete = interpolate_nondiscrete super(Poisson, self).__init__( dtype=dtype, reparameterization_type=reparameterization.NOT_REPARAMETERIZED, validate_args=validate_args, allow_nan_stats=allow_nan_stats, parameters=parameters, name=name)
def __init__(self, df, kernel, index_points=None, mean_fn=None, observation_noise_variance=0., marginal_fn=None, jitter=1e-6, validate_args=False, allow_nan_stats=False, name='StudentTProcess'): """Instantiate a StudentTProcess Distribution. Args: df: Positive Floating-point `Tensor` representing the degrees of freedom. Must be greater than 2. kernel: `PositiveSemidefiniteKernel`-like instance representing the TP's covariance function. index_points: `float` `Tensor` representing finite (batch of) vector(s) of points in the index set over which the TP is defined. Shape has the form `[b1, ..., bB, e, f1, ..., fF]` where `F` is the number of feature dimensions and must equal `kernel.feature_ndims` and `e` is the number (size) of index points in each batch. Ultimately this distribution corresponds to a `e`-dimensional multivariate Student's T. The batch shape must be broadcastable with `kernel.batch_shape` and any batch dims yielded by `mean_fn`. mean_fn: Python `callable` that acts on `index_points` to produce a (batch of) vector(s) of mean values at `index_points`. Takes a `Tensor` of shape `[b1, ..., bB, f1, ..., fF]` and returns a `Tensor` whose shape is broadcastable with `[b1, ..., bB]`. Default value: `None` implies constant zero function. observation_noise_variance: `float` `Tensor` representing (batch of) scalar variance(s) of the noise in the Normal likelihood distribution of the model. If batched, the batch shape must be broadcastable with the shapes of all other batched parameters (`kernel.batch_shape`, `index_points`, etc.). Default value: `0.` marginal_fn: A Python callable that takes a location, covariance matrix, optional `validate_args`, `allow_nan_stats` and `name` arguments, and returns a multivariate normal subclass of `tfd.Distribution`. Default value: `None`, in which case a Cholesky-factorizing function is is created using `make_cholesky_factorizing_marginal_fn` and the `jitter` argument. jitter: `float` scalar `Tensor` added to the diagonal of the covariance matrix to ensure positive definiteness of the covariance matrix. Default value: `1e-6`. validate_args: Python `bool`, default `False`. When `True` distribution parameters are checked for validity despite possibly degrading runtime performance. When `False` invalid inputs may silently render incorrect outputs. Default value: `False`. allow_nan_stats: Python `bool`, default `True`. When `True`, statistics (e.g., mean, mode, variance) use the value "`NaN`" to indicate the result is undefined. When `False`, an exception is raised if one or more of the statistic's batch members are undefined. Default value: `False`. name: Python `str` name prefixed to Ops created by this class. Default value: "StudentTProcess". Raises: ValueError: if `mean_fn` is not `None` and is not callable. """ parameters = dict(locals()) with tf.name_scope(name) as name: dtype = dtype_util.common_dtype( [df, index_points, observation_noise_variance, jitter], tf.float32) df = tensor_util.convert_nonref_to_tensor(df, dtype=dtype, name='df') observation_noise_variance = tensor_util.convert_nonref_to_tensor( observation_noise_variance, dtype=dtype, name='observation_noise_variance') index_points = tensor_util.convert_nonref_to_tensor( index_points, dtype=dtype, name='index_points') jitter = tensor_util.convert_nonref_to_tensor(jitter, dtype=dtype, name='jitter') self._kernel = kernel self._index_points = index_points # Default to a constant zero function, borrowing the dtype from # index_points to ensure consistency. if mean_fn is None: mean_fn = lambda x: tf.zeros([1], dtype=dtype) else: if not callable(mean_fn): raise ValueError('`mean_fn` must be a Python callable') self._df = df self._observation_noise_variance = observation_noise_variance self._mean_fn = mean_fn self._jitter = jitter if marginal_fn is None: self._marginal_fn = make_cholesky_factored_marginal_fn(jitter) else: self._marginal_fn = marginal_fn with tf.name_scope('init'): super(StudentTProcess, self).__init__( dtype=dtype, reparameterization_type=reparameterization. FULLY_REPARAMETERIZED, validate_args=validate_args, allow_nan_stats=allow_nan_stats, parameters=parameters, name=name)
def __init__(self, loc, scale, skewness=None, tailweight=None, distribution=None, validate_args=False, allow_nan_stats=True, name="SinhArcsinh"): """Construct SinhArcsinh distribution on `(-inf, inf)`. Arguments `(loc, scale, skewness, tailweight)` must have broadcastable shape (indexing batch dimensions). They must all have the same `dtype`. Args: loc: Floating-point `Tensor`. scale: `Tensor` of same `dtype` as `loc`. skewness: Skewness parameter. Default is `0.0` (no skew). tailweight: Tailweight parameter. Default is `1.0` (unchanged tailweight) distribution: `tf.Distribution`-like instance. Distribution that is transformed to produce this distribution. Default is `tfd.Normal(0., 1.)`. Must be a scalar-batch, scalar-event distribution. Typically `distribution.reparameterization_type = FULLY_REPARAMETERIZED` or it is a function of non-trainable parameters. WARNING: If you backprop through a `SinhArcsinh` sample and `distribution` is not `FULLY_REPARAMETERIZED` yet is a function of trainable variables, then the gradient will be incorrect! validate_args: Python `bool`, default `False`. When `True` distribution parameters are checked for validity despite possibly degrading runtime performance. When `False` invalid inputs may silently render incorrect outputs. allow_nan_stats: Python `bool`, default `True`. When `True`, statistics (e.g., mean, mode, variance) use the value "`NaN`" to indicate the result is undefined. When `False`, an exception is raised if one or more of the statistic's batch members are undefined. name: Python `str` name prefixed to Ops created by this class. """ parameters = dict(locals()) with tf.name_scope(name) as name: dtype = dtype_util.common_dtype([loc, scale, skewness, tailweight], tf.float32) self._loc = tensor_util.convert_nonref_to_tensor(loc, name="loc", dtype=dtype) self._scale = tensor_util.convert_nonref_to_tensor(scale, name="scale", dtype=dtype) tailweight = 1. if tailweight is None else tailweight has_default_skewness = skewness is None skewness = 0. if has_default_skewness else skewness self._tailweight = tensor_util.convert_nonref_to_tensor( tailweight, name="tailweight", dtype=dtype) self._skewness = tensor_util.convert_nonref_to_tensor( skewness, name="skewness", dtype=dtype) batch_shape = distribution_util.get_broadcast_shape( self._loc, self._scale, self._tailweight, self._skewness) # Recall, with Z a random variable, # Y := loc + scale * F(Z), # F(Z) := Sinh( (Arcsinh(Z) + skewness) * tailweight ) * C # C := 2 / F_0(2) # F_0(Z) := Sinh( Arcsinh(Z) * tailweight ) if distribution is None: distribution = normal.Normal(loc=tf.zeros([], dtype=dtype), scale=tf.ones([], dtype=dtype), allow_nan_stats=allow_nan_stats, validate_args=validate_args) else: asserts = distribution_util.maybe_check_scalar_distribution( distribution, dtype, validate_args) if asserts: self._loc = distribution_util.with_dependencies( asserts, self._loc) # Make the SAS bijector, 'F'. f = sinh_arcsinh_bijector.SinhArcsinh(skewness=self._skewness, tailweight=self._tailweight, validate_args=validate_args) # Make the AffineScalar bijector, Z --> loc + scale * Z (2 / F_0(2)) affine = affine_scalar_bijector.AffineScalar( shift=self._loc, scale=self._scale, validate_args=validate_args) bijector = chain_bijector.Chain([affine, f]) super(SinhArcsinh, self).__init__(distribution=distribution, bijector=bijector, batch_shape=batch_shape, validate_args=validate_args, name=name) self._parameters = parameters