def _calculate_batch_shape(self): """Computes fully defined batch shape for the new distribution.""" all_batch_shapes = [d.batch_shape.as_list() if tensorshape_util.is_fully_defined(d.batch_shape) else d.batch_shape_tensor() for d in self.distributions] original_shape = ps.stack(all_batch_shapes, axis=0) index_mask = ps.cast( ps.one_hot(self._axis, ps.shape(original_shape)[1]), dtype=tf.bool) new_concat_dim = ps.cast( ps.reduce_sum(original_shape, axis=0)[self._axis], dtype=tf.int32) return ps.where(index_mask, new_concat_dim, ps.reduce_max(original_shape, axis=0))
def initial_value_of_masked_time_series(time_series_tensor, broadcast_mask): """Get the first unmasked entry of each time series in the batch. If a batch element has no unmasked entries, the corresponding return value for that element is undefined. Args: time_series_tensor: float `Tensor` of shape `batch_shape + [num_timesteps]`. broadcast_mask: bool `Tensor` of same shape as `time_series`. Returns: initial_values: float `Tensor` of shape `batch_shape`. """ num_timesteps = ps.shape(time_series_tensor)[-1] # Compute the index of the first unmasked entry for each series in the batch. unmasked_negindices = (ps.cast(~broadcast_mask, np.int32) * ps.range(num_timesteps, 0, -1)) first_unmasked_indices = num_timesteps - ps.reduce_max(unmasked_negindices, axis=-1) # Avoid out-of-bounds errors if all indices are masked. safe_unmasked_indices = ps.minimum(first_unmasked_indices, num_timesteps - 1) batch_dims = tensorshape_util.rank(safe_unmasked_indices.shape) if batch_dims is None: raise NotImplementedError( 'Cannot compute initial values of a masked time series with' 'dynamic rank.') # `batch_gather` requires static rank # Extract the initial value for each series in the batch. return tf.squeeze( tf.gather(params=time_series_tensor, indices=safe_unmasked_indices[..., np.newaxis], batch_dims=batch_dims, axis=-1), # Since we've gathered exactly one step from the # `num_timesteps` axis, we can remove that axis entirely. axis=-1)
def _update_forward_min_event_ndims( bij, downstream_quantities, get_forward_min_event_ndims=lambda b: b.forward_min_event_ndims, get_inverse_min_event_ndims=lambda b: b.inverse_min_event_ndims, inverse_event_ndims_fn=lambda b, nd: b.inverse_event_ndims(nd)): """Step backwards through the graph to infer `forward_min_event_ndims`. Args: bij: local tfb.Bijector instance at the current graph node. downstream_quantities: Instance of `MinEventNdimsDownstreamQuantities` namedtuple, containing event_ndims that satisfy the bijector(s) downstream from `bij` in the graph. May be `None` if there are no such bijectors. get_forward_min_event_ndims: callable; may be overridden to swap forward/inverse direction. get_inverse_min_event_ndims: callable; may be overridden to swap forward/inverse direction. inverse_event_ndims_fn: callable; may be overridden to swap forward/inverse direction. Returns: downstream_quantities: Instance of `MinEventNdimsDownstreamQuantities` namedtuple containing event_ndims that satisfy `bij` and all downstream bijectors. """ if downstream_quantities is None: # This is a leaf bijector. return MinEventNdimsInferenceDownstreamQuantities( forward_min_event_ndims=get_forward_min_event_ndims(bij), parts_interact=bij._parts_interact) # pylint: disable=protected-access inverse_min_event_ndims = get_inverse_min_event_ndims(bij) downstream_min_event_ndims = nest_util.coerce_structure( inverse_min_event_ndims, downstream_quantities.forward_min_event_ndims) # Update the min_event_ndims that is a valid input to downstream bijectors # to also be a valid *output* of this bijector, or equivalently, a valid # input to `bij.inverse`. rank_mismatches = tf.nest.flatten( tf.nest.map_structure( lambda dim, min_dim: dim - min_dim, downstream_min_event_ndims, inverse_min_event_ndims)) if downstream_quantities.parts_interact: # If downstream bijectors involve interaction between parts, # then a valid input to the downstream bijectors must augment the # `downstream_min_event_ndims` by the # same rank for every part (otherwise we would induce event shape # broadcasting). Hopefully, this will also avoid event-shape broadcasting # at the current bijector---if not, the composition is invalid, and the call # to `bij.inverse_event_ndims(valid_inverse_min_event_ndims)` below will # raise an exception. maximum_rank_deficiency = -ps.reduce_min([0] + rank_mismatches) valid_inverse_min_event_ndims = tf.nest.map_structure( lambda ndims: maximum_rank_deficiency + ndims, downstream_min_event_ndims) else: if bij._parts_interact: # pylint: disable=protected-access # If this bijector does *not* operate independently on its parts, then a # valid input to `inverse` cannot require event shape broadcasting. That # is, each part must have the same 'excess rank' above the local # inverse_min_event_ndims; we ensure this by construction. maximum_excess_rank = ps.reduce_max([0] + rank_mismatches) valid_inverse_min_event_ndims = tf.nest.map_structure( lambda ndims: maximum_excess_rank + ndims, inverse_min_event_ndims) else: # If all parts are independent, can take the pointwise max event_ndims. valid_inverse_min_event_ndims = tf.nest.map_structure( ps.maximum, downstream_min_event_ndims, inverse_min_event_ndims) return MinEventNdimsInferenceDownstreamQuantities( # Pull the desired output ndims back through the bijector, to get # the ndims of a valid *input*. forward_min_event_ndims=inverse_event_ndims_fn( bij, valid_inverse_min_event_ndims), parts_interact=( downstream_quantities.parts_interact or bij._parts_interact)) # pylint: disable=protected-access
def __init__(self, skewness, tailweight, loc, scale, validate_args=False, allow_nan_stats=True, name=None): """Construct Johnson's SU distributions. The distributions have shape parameteres `tailweight` and `skewness`, mean `loc`, and scale `scale`. The parameters `tailweight`, `skewness`, `loc`, and `scale` must be shaped in a way that supports broadcasting (e.g. `skewness + tailweight + loc + scale` is a valid operation). Args: skewness: Floating-point `Tensor`. Skewness of the distribution(s). tailweight: Floating-point `Tensor`. Tail weight of the distribution(s). `tailweight` must contain only positive values. loc: Floating-point `Tensor`. The mean(s) of the distribution(s). scale: Floating-point `Tensor`. The scaling factor(s) for the distribution(s). Note that `scale` is not technically the standard deviation of this distribution but has semantics more similar to standard deviation than variance. 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 any of skewness, tailweight, loc and scale are different dtypes. """ parameters = dict(locals()) with tf.name_scope(name or 'JohnsonSU') as name: dtype = dtype_util.common_dtype([skewness, tailweight, loc, scale], tf.float32) self._skewness = tensor_util.convert_nonref_to_tensor( skewness, name='skewness', dtype=dtype) self._tailweight = tensor_util.convert_nonref_to_tensor( tailweight, name='tailweight', dtype=dtype) 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) norm_shift = invert_bijector.Invert( shift_bijector.Shift(shift=self._skewness, validate_args=validate_args)) norm_scale = invert_bijector.Invert( scale_bijector.Scale(scale=self._tailweight, validate_args=validate_args)) sinh = sinh_bijector.Sinh(validate_args=validate_args) scale = scale_bijector.Scale(scale=self._scale, validate_args=validate_args) shift = shift_bijector.Shift(shift=self._loc, validate_args=validate_args) bijector = shift(scale(sinh(norm_scale(norm_shift)))) batch_rank = ps.reduce_max([ distribution_util.prefer_static_rank(x) for x in (self._skewness, self._tailweight, self._loc, self._scale) ]) super(JohnsonSU, self).__init__( # TODO(b/160730249): Make `loc` a scalar `0.` and remove overridden # `batch_shape` and `batch_shape_tensor` when # TransformedDistribution's bijector can modify its `batch_shape`. distribution=normal.Normal(loc=tf.zeros(ps.ones( batch_rank, tf.int32), dtype=dtype), scale=tf.ones([], dtype=dtype), validate_args=validate_args, allow_nan_stats=allow_nan_stats), bijector=bijector, validate_args=validate_args, parameters=parameters, name=name)