def testBijector(self): low = np.array([[-3.], [0.], [5.]]).astype(np.float32) high = 12. bijector = tfb.Sigmoid(low=low, high=high, validate_args=True) equivalent_bijector = tfb.Chain( [tfb.Shift(shift=low), tfb.Scale(scale=high - low), tfb.Sigmoid()]) x = [[[1., 2., -5., -0.3]]] y = self.evaluate(equivalent_bijector.forward(x)) self.assertAllClose(y, self.evaluate(bijector.forward(x))) self.assertAllClose(x, self.evaluate(bijector.inverse(y)[..., :1, :]), rtol=1e-5) self.assertAllClose( self.evaluate( equivalent_bijector.inverse_log_det_jacobian(y, event_ndims=1)), self.evaluate(bijector.inverse_log_det_jacobian(y, event_ndims=1)), rtol=1e-5) self.assertAllClose( self.evaluate( equivalent_bijector.forward_log_det_jacobian(x, event_ndims=1)), self.evaluate(bijector.forward_log_det_jacobian(x, event_ndims=1)))
def testBijector(self): with self.cached_session(): self.assertEqual("sigmoid", tfb.Sigmoid().name) x = np.linspace(-10., 10., 100).reshape([2, 5, 10]).astype(np.float32) y = special.expit(x) ildj = -np.log(y) - np.log1p(-y) bijector = tfb.Sigmoid() self.assertAllClose(y, self.evaluate(bijector.forward(x)), atol=0., rtol=1e-2) self.assertAllClose(x, self.evaluate(bijector.inverse(y)), atol=0., rtol=1e-4) self.assertAllClose(ildj, self.evaluate( bijector.inverse_log_det_jacobian( y, event_ndims=0)), atol=0., rtol=1e-6) self.assertAllClose(-ildj, self.evaluate( bijector.forward_log_det_jacobian( x, event_ndims=0)), atol=0., rtol=1e-4)
def make_asvi_trainable_variables(prior): """Generates parameter dictionaries given a prior distribution and list.""" with tf.name_scope('make_asvi_trainable_variables'): param_dicts = [] prior_dists = prior._get_single_sample_distributions() # pylint: disable=protected-access for dist in prior_dists: actual_dist = dist.distribution if isinstance(dist, Root) else dist dist_params = actual_dist.parameters new_params_dict = {} # Build trainable ASVI representation for each distribution's parameters. for param, value in dist_params.items(): if param in _NON_STATISTICAL_PARAMS or value is None: continue new_params_dict[param] = ASVIParameters( prior_weight=tfp.util.TransformedVariable( 0.5, bijector=tfb.Sigmoid(), name='prior_weight/{}/{}'.format(dist.name, param)), mean_field_parameter=tfp.util.TransformedVariable( 0.5, bijector=parameter_constraints.constraint_for(param), name='mean_field_parameter/{}/{}'.format(dist.name, param))) param_dicts.append(new_params_dict) return param_dicts
def _make_asvi_trainable_variables(prior, mean_field=False, initial_prior_weight=0.5): """Generates parameter dictionaries given a prior distribution and list.""" with tf.name_scope('make_asvi_trainable_variables'): param_dicts = [] prior_dists = prior._get_single_sample_distributions() # pylint: disable=protected-access for dist in prior_dists: original_dist = dist.distribution if isinstance(dist, Root) else dist substituted_dist = _as_trainable_family(original_dist) # Grab the base distribution if it exists try: actual_dist = substituted_dist.distribution except AttributeError: actual_dist = substituted_dist new_params_dict = {} # Build trainable ASVI representation for each distribution's parameters. parameter_properties = actual_dist.parameter_properties( dtype=actual_dist.dtype) sample_shape = tf.concat( [dist.batch_shape_tensor(), dist.event_shape_tensor()], axis=0) for param, value in actual_dist.parameters.items(): if param in (_NON_STATISTICAL_PARAMS + _NON_TRAINABLE_PARAMS) or value is None: continue try: bijector = parameter_properties[ param].default_constraining_bijector_fn() except NotImplementedError: bijector = tfb.Identity() unconstrained_ones = tf.ones( shape=bijector.inverse_event_shape_tensor( parameter_properties[param].shape_fn( sample_shape=sample_shape)), dtype=actual_dist.dtype) if mean_field: new_params_dict[param] = ASVIParameters( prior_weight=None, mean_field_parameter=tfp_util.TransformedVariable( value, bijector=bijector, name='mean_field_parameter/{}/{}'.format(dist.name, param))) else: new_params_dict[param] = ASVIParameters( prior_weight=tfp_util.TransformedVariable( initial_prior_weight * unconstrained_ones, bijector=tfb.Sigmoid(), name='prior_weight/{}/{}'.format(dist.name, param)), mean_field_parameter=tfp_util.TransformedVariable( value, bijector=bijector, name='mean_field_parameter/{}/{}'.format(dist.name, param))) param_dicts.append(new_params_dict) return param_dicts
def testAssertHighGtLow(self): low = np.array([1., 1., 1.], dtype=np.float32) high = np.array([1., 2., 3.], dtype=np.float32) with self.assertRaisesOpError('not defined when `low` >= `high`'): bijector = tfb.Sigmoid(low=low, high=high, validate_args=True) self.evaluate(bijector.forward(3.))
def testSigmoidBetaTargetConservation(self): # Not inverting the sigmoid bijector makes a kooky distribution, but nuts # should still conserve it (with a smaller step size). sigmoid_beta_dist = tfb.Identity(tfb.Sigmoid())( tfd.Beta(concentration0=1., concentration1=2.)) self.evaluate(assert_univariate_target_conservation( self, sigmoid_beta_dist, step_size=0.02))
def testLogitBetaTargetConservation(self): logit_beta_dist = tfb.Invert(tfb.Sigmoid())(tfd.Beta( concentration0=1., concentration1=2.)) self.evaluate( assert_univariate_target_conservation(self, logit_beta_dist, step_size=0.2))
def make_asvi_trainable_variables(prior): """Generates parameter dictionaries given a prior distribution and list.""" param_dicts = [] prior_dists = prior._get_single_sample_distributions() # pylint: disable=protected-access d = 0 for dist in prior_dists: actual_dist = dist.distribution if isinstance(dist, Root) else dist dist_params = actual_dist.parameters new_params_dict = {} # make and store mean_field_parameters & prior_weights in new_params_dict for param, value in dist_params.items(): if param in _NON_STATISTICAL_PARAMS or value is None: new_params_dict[param] = value else: # TODO(kateslin): Change dictionary naming scheme to use namedtuples. prior_weight_name = 'prior_weight_{}'.format(d) mean_field_parameter_name = 'mean_field_parameter_{}'.format(d) prior_weight = tfp.util.TransformedVariable( 0.5, bijector=tfb.Sigmoid(), name=prior_weight_name) new_params_dict['{}_prior_weight'.format(param)] = prior_weight mean_field_parameter_constraint_bijector = ( parameter_constraints.constraint_for(param)) mean_field_parameter = tfp.util.TransformedVariable( 0.5, mean_field_parameter_constraint_bijector, name=mean_field_parameter_name) new_params_dict['{}_mean_field_parameter'.format( param)] = mean_field_parameter d += 1 param_dicts.append(new_params_dict) return param_dicts
def testBijectiveAndFinite(self): x = np.linspace(-100., 100., 100).astype(np.float32) eps = 1e-3 y = np.linspace(eps, 1. - eps, 100).astype(np.float32) bijector_test_util.assert_bijective_and_finite( tfb.Sigmoid(), x, y, eval_func=self.evaluate, event_ndims=0, atol=0., rtol=1e-4)
def testMatchWithAffineTransform(self): direct_bj = tfb.Tanh() indirect_bj = tfb.Chain([ tfb.AffineScalar(shift=tf.to_double(-1.0), scale=tf.to_double(2.0)), tfb.Sigmoid(), tfb.AffineScalar(scale=tf.to_double(2.0)) ]) x = np.linspace(-3.0, 3.0, 100) y = np.tanh(x) self.assertAllClose(self.evaluate(direct_bj.forward(x)), self.evaluate(indirect_bj.forward(x))) self.assertAllClose(self.evaluate(direct_bj.inverse(y)), self.evaluate(indirect_bj.inverse(y))) self.assertAllClose( self.evaluate(direct_bj.inverse_log_det_jacobian(y, event_ndims=0)), self.evaluate( indirect_bj.inverse_log_det_jacobian(y, event_ndims=0))) self.assertAllClose( self.evaluate(direct_bj.forward_log_det_jacobian(x, event_ndims=0)), self.evaluate( indirect_bj.forward_log_det_jacobian(x, event_ndims=0)))
def testMatchWithAffineTransform(self): direct_bj = tfb.Tanh() indirect_bj = tfb.Chain([ tfb.Shift(tf.cast(-1.0, dtype=tf.float64)), tfb.Scale(tf.cast(2.0, dtype=tf.float64)), tfb.Sigmoid(), tfb.Scale(tf.cast(2.0, dtype=tf.float64)) ]) x = np.linspace(-3.0, 3.0, 100) y = np.tanh(x) self.assertAllClose(self.evaluate(direct_bj.forward(x)), self.evaluate(indirect_bj.forward(x))) self.assertAllClose(self.evaluate(direct_bj.inverse(y)), self.evaluate(indirect_bj.inverse(y))) self.assertAllClose( self.evaluate(direct_bj.inverse_log_det_jacobian(y, event_ndims=0)), self.evaluate( indirect_bj.inverse_log_det_jacobian(y, event_ndims=0))) self.assertAllClose( self.evaluate(direct_bj.forward_log_det_jacobian(x, event_ndims=0)), self.evaluate( indirect_bj.forward_log_det_jacobian(x, event_ndims=0)))
def testScalarCongruency(self): low = -2. high = 5. bijector = tfb.Sigmoid(low=low, high=high, validate_args=True) bijector_test_util.assert_scalar_congruency(bijector, lower_x=-5., upper_x=3.5, eval_func=self.evaluate, rtol=0.05)
def testNumericallySuperiorToEquivalentChain(self): x = np.array([-5., 3., 17., 23.]).astype(np.float32) low = -0.08587775 high = 0.12498104 bijector = tfb.Sigmoid(low=low, high=high, validate_args=True) equivalent_bijector = tfb.Chain( [tfb.Shift(shift=low), tfb.Scale(scale=high - low), tfb.Sigmoid()]) self.assertAllLessEqual(self.evaluate(bijector.forward(x)), high) # The mathematically equivalent `Chain` bijector can return values greater # than the intended upper bound of `high`. self.assertTrue( (self.evaluate(equivalent_bijector.forward(x)) > high).any())
def testQuantile(self): logit_normal = self._cls()(distribution=tfd.Normal(loc=0., scale=1.), bijector=tfb.Sigmoid(), validate_args=True) grid = [0., 0.25, 0.5, 0.75, 1.] q = logit_normal.quantile(grid) cdf = logit_normal.cdf(q) cdf_ = self.evaluate(cdf) self.assertAllClose(grid, cdf_, rtol=1e-6, atol=0.)
def testCopyOverride(self): sigmoid = tfb.Sigmoid(low=-1., high=2., validate_args=True) self.assertEqual(sigmoid.parameters, sigmoid.copy().parameters) unused_sigmoid_copy = sigmoid.copy(validate_args=False) base_params = sigmoid.parameters.copy() copy_params = sigmoid.copy(validate_args=False).parameters.copy() self.assertNotEqual( base_params.pop('validate_args'), copy_params.pop('validate_args')) self.assertEqual(base_params, copy_params)
def __init__(self, temperature, logits=None, probs=None, validate_args=False, allow_nan_stats=True, name="RelaxedBernoulli"): """Construct RelaxedBernoulli distributions. Args: temperature: An 0-D `Tensor`, representing the temperature of a set of RelaxedBernoulli distributions. The temperature should be positive. logits: An N-D `Tensor` representing the log-odds of a positive event. Each entry in the `Tensor` parametrizes an independent RelaxedBernoulli distribution where the probability of an event is sigmoid(logits). Only one of `logits` or `probs` should be passed in. probs: An N-D `Tensor` representing the probability of a positive event. Each entry in the `Tensor` parameterizes an independent Bernoulli distribution. Only one of `logits` or `probs` should be passed in. 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 both `probs` and `logits` are passed, or if neither. """ parameters = dict(locals()) with tf.name_scope(name, values=[logits, probs, temperature]) as name: dtype = dtype_util.common_dtype([logits, probs, temperature], tf.float32) self._temperature = tf.convert_to_tensor( temperature, name="temperature", dtype=dtype) if validate_args: with tf.control_dependencies([tf.assert_positive(temperature)]): self._temperature = tf.identity(self._temperature) self._logits, self._probs = distribution_util.get_logits_and_probs( logits=logits, probs=probs, validate_args=validate_args, dtype=dtype) super(RelaxedBernoulli, self).__init__( distribution=logistic.Logistic( self._logits / self._temperature, 1. / self._temperature, validate_args=validate_args, allow_nan_stats=allow_nan_stats, name=name + "/Logistic"), bijector=bijectors.Sigmoid(validate_args=validate_args), validate_args=validate_args, name=name) self._parameters = parameters
def __init__(self, centered_returns, name='stochastic_volatility', pretty_name='Stochastic Volatility'): """Construct the stochastic volatility model. Args: centered_returns: Float `Tensor` of shape `[num_timesteps]` giving the mean-adjusted return (change in asset price, minus the average change) observed at each step. name: Python `str` name prefixed to Ops created by this class. pretty_name: A Python `str`. The pretty name of this model. """ with tf.name_scope(name): num_timesteps = ps.size0(centered_returns) if tf.is_tensor(num_timesteps): raise ValueError( 'Returns series length must be static, but saw ' 'shape {}.'.format(centered_returns.shape)) self._prior_dist = tfd.JointDistributionCoroutine( functools.partial(stochastic_volatility_prior_fn, num_timesteps=num_timesteps)) self._log_likelihood_fn = functools.partial( stochastic_volatility_log_likelihood_fn, centered_returns=centered_returns) def _ext_identity(params): res = collections.OrderedDict() res['persistence_of_volatility'] = params[0] res['mean_log_volatility'] = params[1] res['white_noise_shock_scale'] = params[2] res['log_volatility'] = tf.stack(params[3], axis=-1) return res sample_transformations = { 'identity': model.Model.SampleTransformation( fn=_ext_identity, pretty_name='Identity', ) } super(StochasticVolatility, self).__init__( default_event_space_bijector=(tfb.Sigmoid( -1., 1.), tfb.Identity(), tfb.Softplus()) + ((tfb.Identity(), ) * num_timesteps, ), event_shape=self._prior_dist.event_shape, dtype=self._prior_dist.dtype, name=name, pretty_name=pretty_name, sample_transformations=sample_transformations, )
def testCopyExtraArgs(self): # Note: we cannot easily test all bijectors since each requires # different initialization arguments. We therefore spot test a few. sigmoid = tfb.Sigmoid(low=-1., high=2., validate_args=True) self.assertEqual(sigmoid.parameters, sigmoid.copy().parameters) chain = tfb.Chain( [ tfb.Softplus(hinge_softness=[1., 2.], validate_args=True), tfb.MatrixInverseTriL(validate_args=True) ], validate_args=True) self.assertEqual(chain.parameters, chain.copy().parameters)
def testBijectiveAndFinite(self): with self.test_session(): x = np.linspace(-7., 7., 100).astype(np.float32) eps = 1e-3 y = np.linspace(eps, 1. - eps, 100).astype(np.float32) assert_bijective_and_finite(tfb.Sigmoid(), x, y, event_ndims=0, atol=0., rtol=1e-4)
def testBijectiveAndFinite(self): low = -5. high = 8. bijector = tfb.Sigmoid(low=low, high=high, validate_args=True) x = np.linspace(-10, 10, num=100).astype(np.float32) eps = 1e-6 y = np.linspace(low + eps, high - eps, num=100).astype(np.float32) bijector_test_util.assert_bijective_and_finite(bijector, x, y, eval_func=self.evaluate, event_ndims=0)
def testEdgeCaseRequiringClipping(self): np.set_printoptions(floatmode='unique', precision=None) lo = np.float32(0.010489981) hi = test_util.floats_near(0.010499111, 100, dtype=np.float32)[:, np.newaxis] self.assertAllEqual([100, 1], hi.shape) xs = test_util.floats_near(9.814646, 100, dtype=np.float32) bijector = tfb.Sigmoid(low=lo, high=hi, validate_args=True) answers = bijector.forward(xs) self.assertAllEqual([100, 100], answers.shape) for ans1, hi1 in zip(self.evaluate(answers), hi): self.assertAllLessEqual(ans1, hi1)
def test_support_works_correctly_with_HMC(self): num_results = 2000 target = tfd.Beta(concentration1=self.dtype(1.), concentration0=self.dtype(10.)) transformed_hmc = tfp.mcmc.TransformedTransitionKernel( inner_kernel=tfp.mcmc.HamiltonianMonteCarlo( target_log_prob_fn=tf.function(target.log_prob, autograph=False), step_size=1.64, num_leapfrog_steps=2, seed=_maybe_seed(55)), bijector=tfb.Sigmoid()) # Recall, tfp.mcmc.sample_chain calls # transformed_hmc.bootstrap_results too. states, kernel_results = tfp.mcmc.sample_chain( num_results=num_results, # The initial state is used by inner_kernel.bootstrap_results. # Note the input is *after* bijector.forward. current_state=self.dtype(0.25), kernel=transformed_hmc, num_burnin_steps=200, num_steps_between_results=1, parallel_iterations=1) self.assertEqual(num_results, tf.compat.dimension_value(states.shape[0])) sample_mean = tf.reduce_mean(states, axis=0) sample_var = tf.reduce_mean(tf.math.squared_difference( states, sample_mean), axis=0) [ sample_mean_, sample_var_, is_accepted_, true_mean_, true_var_, ] = self.evaluate([ sample_mean, sample_var, kernel_results.inner_results.is_accepted, target.mean(), target.variance(), ]) self.assertAllClose(true_mean_, sample_mean_, atol=0.06, rtol=0.) self.assertAllClose(true_var_, sample_var_, atol=0.01, rtol=0.1) self.assertNear(0.6, is_accepted_.mean(), err=0.05)
def test_support_works_correctly_with_MALA(self): num_results = 2000 target = tfd.Beta(concentration1=self.dtype(1.), concentration0=self.dtype(10.)) transformed_mala = tfp.mcmc.TransformedTransitionKernel( inner_kernel=tfp.mcmc.MetropolisAdjustedLangevinAlgorithm( target_log_prob_fn=tf.function(target.log_prob, autograph=False), step_size=1., seed=_maybe_seed(test_util.test_seed())), bijector=tfb.Sigmoid()) # Recall, tfp.mcmc.sample_chain calls # transformed_hmc.bootstrap_results too. states, _ = tfp.mcmc.sample_chain( num_results=num_results, # The initial state is used by inner_kernel.bootstrap_results. # Note the input is *after* bijector.forward. current_state=self.dtype(0.25), kernel=transformed_mala, num_burnin_steps=200, num_steps_between_results=1, parallel_iterations=1) self.assertEqual(num_results, tf.compat.dimension_value(states.shape[0])) sample_mean = tf.reduce_mean(states, axis=0) sample_var = tf.reduce_mean(tf.math.squared_difference( states, sample_mean), axis=0) [ sample_mean_, sample_var_, true_mean_, true_var_, ] = self.evaluate([ sample_mean, sample_var, target.mean(), target.variance(), ]) self.assertAllClose(true_mean_, sample_mean_, atol=0.06, rtol=0.) self.assertAllClose(true_var_, sample_var_, atol=0.01, rtol=0.1)
def test_support_works_correctly_with_rwm(self): num_results = 500 target = tfd.Beta( concentration1=self.dtype(1.), concentration0=self.dtype(10.)) transformed_rwm = tfp.mcmc.TransformedTransitionKernel( inner_kernel=tfp.mcmc.RandomWalkMetropolis( target_log_prob_fn=tf.function(target.log_prob, autograph=False), new_state_fn=tfp.mcmc.random_walk_normal_fn(scale=1.5)), bijector=tfb.Sigmoid()) # Recall, tfp.mcmc.sample_chain calls # transformed_hmc.bootstrap_results too. states = tfp.mcmc.sample_chain( num_results=num_results, # The initial state is used by inner_kernel.bootstrap_results. # Note the input is *after* bijector.forward. current_state=self.dtype(0.25), kernel=transformed_rwm, num_burnin_steps=200, num_steps_between_results=1, trace_fn=None, seed=test_util.test_seed()) self.assertEqual(num_results, tf.compat.dimension_value(states.shape[0])) sample_mean = tf.reduce_mean(states, axis=0) sample_var = tf.reduce_mean( tf.math.squared_difference(states, sample_mean), axis=0) [ sample_mean_, sample_var_, true_mean_, true_var_, ] = self.evaluate([ sample_mean, sample_var, target.mean(), target.variance(), ]) self.assertAllClose(true_mean_, sample_mean_, atol=0.15, rtol=0.) self.assertAllClose(true_var_, sample_var_, atol=0.03, rtol=0.2)
def __init__(self, level_scale_prior=None, slope_mean_prior=None, slope_scale_prior=None, autoregressive_coef_prior=None, initial_level_prior=None, initial_slope_prior=None, observed_time_series=None, constrain_ar_coef_stationary=True, constrain_ar_coef_positive=False, name=None): """Specify a semi-local linear trend model. Args: level_scale_prior: optional `tfd.Distribution` instance specifying a prior on the `level_scale` parameter. If `None`, a heuristic default prior is constructed based on the provided `observed_time_series`. Default value: `None`. slope_mean_prior: optional `tfd.Distribution` instance specifying a prior on the `slope_mean` parameter. If `None`, a heuristic default prior is constructed based on the provided `observed_time_series`. Default value: `None`. slope_scale_prior: optional `tfd.Distribution` instance specifying a prior on the `slope_scale` parameter. If `None`, a heuristic default prior is constructed based on the provided `observed_time_series`. Default value: `None`. autoregressive_coef_prior: optional `tfd.Distribution` instance specifying a prior on the `autoregressive_coef` parameter. If `None`, the default prior is a standard `Normal(0., 1.)`. Note that the prior may be implicitly truncated by `constrain_ar_coef_stationary` and/or `constrain_ar_coef_positive`. Default value: `None`. initial_level_prior: optional `tfd.Distribution` instance specifying a prior on the initial level. If `None`, a heuristic default prior is constructed based on the provided `observed_time_series`. Default value: `None`. initial_slope_prior: optional `tfd.Distribution` instance specifying a prior on the initial slope. If `None`, a heuristic default prior is constructed based on the provided `observed_time_series`. Default value: `None`. observed_time_series: optional `float` `Tensor` of shape `batch_shape + [T, 1]` (omitting the trailing unit dimension is also supported when `T > 1`), specifying an observed time series. Any priors not explicitly set will be given default values according to the scale of the observed time series (or batch of time series). May optionally be an instance of `tfp.sts.MaskedTimeSeries`, which includes a mask `Tensor` to specify timesteps with missing observations. Default value: `None`. constrain_ar_coef_stationary: if `True`, perform inference using a parameterization that restricts `autoregressive_coef` to the interval `(-1, 1)`, or `(0, 1)` if `force_positive_ar_coef` is also `True`, corresponding to stationary processes. This will implicitly truncates the support of `autoregressive_coef_prior`. Default value: `True`. constrain_ar_coef_positive: if `True`, perform inference using a parameterization that restricts `autoregressive_coef` to be positive, or in `(0, 1)` if `constrain_ar_coef_stationary` is also `True`. This will implicitly truncate the support of `autoregressive_coef_prior`. Default value: `False`. name: the name of this model component. Default value: 'SemiLocalLinearTrend'. """ with tf.name_scope(name or 'SemiLocalLinearTrend') as name: if observed_time_series is not None: _, observed_stddev, observed_initial = sts_util.empirical_statistics( observed_time_series) else: observed_stddev, observed_initial = 1., 0. # Heuristic default priors. Overriding these may dramatically # change inference performance and results. if level_scale_prior is None: level_scale_prior = tfd.LogNormal(loc=tf.math.log( .01 * observed_stddev), scale=2.) if slope_mean_prior is None: slope_mean_prior = tfd.Normal(loc=0., scale=observed_stddev) if slope_scale_prior is None: slope_scale_prior = tfd.LogNormal(loc=tf.math.log( .01 * observed_stddev), scale=2.) if autoregressive_coef_prior is None: autoregressive_coef_prior = tfd.Normal( loc=0., scale=tf.ones_like(observed_initial)) if initial_level_prior is None: initial_level_prior = tfd.Normal( loc=observed_initial, scale=tf.abs(observed_initial) + observed_stddev) if initial_slope_prior is None: initial_slope_prior = tfd.Normal(loc=0., scale=observed_stddev) self._initial_state_prior = tfd.MultivariateNormalDiag( loc=tf.stack( [initial_level_prior.mean(), initial_slope_prior.mean()], axis=-1), scale_diag=tf.stack([ initial_level_prior.stddev(), initial_slope_prior.stddev() ], axis=-1)) # Constrain the support of the autoregressive coefficient. if constrain_ar_coef_stationary and constrain_ar_coef_positive: autoregressive_coef_bijector = tfb.Sigmoid( ) # support in (0, 1) elif constrain_ar_coef_positive: autoregressive_coef_bijector = tfb.Softplus( ) # support in (0, infty) elif constrain_ar_coef_stationary: autoregressive_coef_bijector = tfb.Tanh() # support in (-1, 1) else: autoregressive_coef_bijector = tfb.Identity() # unconstrained stddev_preconditioner = tfb.Scale(scale=observed_stddev) scaled_softplus = tfb.Chain( [stddev_preconditioner, tfb.Softplus()]) super(SemiLocalLinearTrend, self).__init__(parameters=[ Parameter('level_scale', level_scale_prior, scaled_softplus), Parameter('slope_mean', slope_mean_prior, stddev_preconditioner), Parameter('slope_scale', slope_scale_prior, scaled_softplus), Parameter('autoregressive_coef', autoregressive_coef_prior, autoregressive_coef_bijector), ], latent_size=2, name=name)
def testScalarCongruency(self): bijector_test_util.assert_scalar_congruency(tfb.Sigmoid(), lower_x=-7., upper_x=7., eval_func=self.evaluate)
def mk_sigmoid_beta(): beta = tfd.Beta(concentration0=1., concentration1=2.) # Not inverting the sigmoid bijector makes a kooky distribution, but # nuts should still conserve it (with a smaller step size). return tfb.Sigmoid()(beta)
def mk_logit_beta(): beta = tfd.Beta(concentration0=1., concentration1=2.) return tfb.Invert(tfb.Sigmoid())(beta)
def _make_asvi_trainable_variables(prior, mean_field=False, initial_prior_weight=0.5): """Generates parameter dictionaries given a prior distribution and list.""" with tf.name_scope('make_asvi_trainable_variables'): param_dicts = [] prior_dists = prior._get_single_sample_distributions() # pylint: disable=protected-access for dist in prior_dists: original_dist = dist.distribution if isinstance(dist, Root) else dist substituted_dist = _as_trainable_family(original_dist) # Grab the base distribution if it exists try: actual_dist = substituted_dist.distribution except AttributeError: actual_dist = substituted_dist new_params_dict = {} # Build trainable ASVI representation for each distribution's parameters. parameter_properties = actual_dist.parameter_properties( dtype=actual_dist.dtype) if isinstance(original_dist, sample.Sample): posterior_batch_shape = ps.concat([ actual_dist.batch_shape_tensor(), original_dist.sample_shape ], axis=0) else: posterior_batch_shape = actual_dist.batch_shape_tensor() for param, value in actual_dist.parameters.items(): if param in (_NON_STATISTICAL_PARAMS + _NON_TRAINABLE_PARAMS) or value is None: continue actual_event_shape = parameter_properties[param].shape_fn( actual_dist.event_shape_tensor()) try: bijector = parameter_properties[ param].default_constraining_bijector_fn() except NotImplementedError: bijector = tfb.Identity() if mean_field: prior_weight = None else: unconstrained_ones = tf.ones(shape=ps.concat([ posterior_batch_shape, bijector.inverse_event_shape_tensor(actual_event_shape) ], axis=0), dtype=actual_dist.dtype) prior_weight = tfp_util.TransformedVariable( initial_prior_weight * unconstrained_ones, bijector=tfb.Sigmoid(), name='prior_weight/{}/{}'.format(dist.name, param)) # If the prior distribution was a tfd.Sample wrapping a base # distribution, we want to give every single sample in the prior its # own lambda and alpha value (rather than having a single lambda and # alpha). if isinstance(original_dist, sample.Sample): value = tf.reshape( value, ps.concat([ actual_dist.batch_shape_tensor(), ps.ones( ps.rank_from_shape( original_dist.sample_shape)), actual_event_shape ], axis=0)) value = tf.broadcast_to( value, ps.concat([posterior_batch_shape, actual_event_shape], axis=0)) new_params_dict[param] = ASVIParameters( prior_weight=prior_weight, mean_field_parameter=tfp_util.TransformedVariable( value, bijector=bijector, name='mean_field_parameter/{}/{}'.format( dist.name, param))) param_dicts.append(new_params_dict) return param_dicts
def fn(x): return tf.math.log(tfb.Sigmoid().forward(x))