def test_dict_sample_log_prob(self): # pylint: disable=bad-whitespace d = tfd.JointDistributionNamed(dict( e=tfd.Independent(tfd.Exponential(rate=[100, 120]), 1), scale=lambda e: tfd.Gamma(concentration=e[..., 0], rate=e[..., 1]), loc=tfd.Normal(loc=0, scale=2.), m=tfd.Normal, x=lambda m: tfd.Sample(tfd.Bernoulli(logits=m), 12)), validate_args=True) # pylint: enable=bad-whitespace self.assertEqual(( ('e', ()), ('scale', ('e', )), ('loc', ()), ('m', ('loc', 'scale')), ('x', ('m', )), ), d._resolve_graph()) xs = d.sample(seed=test_util.test_seed()) self.assertLen(xs, 5) # We'll verify the shapes work as intended when we plumb these back into the # respective log_probs. ds, _ = d.sample_distributions(value=xs) self.assertLen(ds, 5) self.assertIsInstance(ds['e'], tfd.Independent) self.assertIsInstance(ds['scale'], tfd.Gamma) self.assertIsInstance(ds['loc'], tfd.Normal) self.assertIsInstance(ds['m'], tfd.Normal) self.assertIsInstance(ds['x'], tfd.Sample) # Static properties. self.assertAllEqual( { 'e': tf.float32, 'scale': tf.float32, 'loc': tf.float32, 'm': tf.float32, 'x': tf.int32 }, d.dtype) batch_shape_tensor_, event_shape_tensor_ = self.evaluate( [d.batch_shape_tensor(), d.event_shape_tensor()]) expected_batch_shape = { 'e': [], 'scale': [], 'loc': [], 'm': [], 'x': [] } batch_tensorshape = d.batch_shape for k in expected_batch_shape: self.assertAllEqual(expected_batch_shape[k], batch_tensorshape[k]) self.assertAllEqual(expected_batch_shape[k], batch_shape_tensor_[k]) expected_event_shape = { 'e': [2], 'scale': [], 'loc': [], 'm': [], 'x': [12] } event_tensorshape = d.event_shape for k in expected_event_shape: self.assertAllEqual(expected_event_shape[k], event_tensorshape[k]) self.assertAllEqual(expected_event_shape[k], event_shape_tensor_[k]) expected_jlp = sum(ds[k].log_prob(xs[k]) for k in ds.keys()) actual_jlp = d.log_prob(xs) self.assertAllClose(*self.evaluate([expected_jlp, actual_jlp]), atol=0., rtol=1e-4)
def setUp(self): self.dtype = np.float32 self.model = tfp.glm.GammaExp() one = np.array(1, self.dtype) self.expected = tfp.glm.CustomExponentialFamily( lambda mu: tfd.Gamma(concentration=one, rate=1. / mu), tf.exp)
def test_namedtuple_sample_log_prob(self): Model = collections.namedtuple('Model', ['e', 'scale', 'loc', 'm', 'x']) # pylint: disable=invalid-name # pylint: disable=bad-whitespace model = Model( e=tfd.Independent(tfd.Exponential(rate=[100, 120]), 1), scale=lambda e: tfd.Gamma(concentration=e[..., 0], rate=e[..., 1]), loc=tfd.Normal(loc=0, scale=2.), m=tfd.Normal, x=lambda m: tfd.Sample(tfd.Bernoulli(logits=m), 12)) # pylint: enable=bad-whitespace d = tfd.JointDistributionNamed(model, validate_args=True) self.assertEqual(( ('e', ()), ('scale', ('e', )), ('loc', ()), ('m', ('loc', 'scale')), ('x', ('m', )), ), d._resolve_graph()) xs = d.sample(seed=test_util.test_seed()) self.assertLen(xs, 5) # We'll verify the shapes work as intended when we plumb these back into the # respective log_probs. ds, _ = d.sample_distributions(value=xs) self.assertLen(ds, 5) self.assertIsInstance(ds.e, tfd.Independent) self.assertIsInstance(ds.scale, tfd.Gamma) self.assertIsInstance(ds.loc, tfd.Normal) self.assertIsInstance(ds.m, tfd.Normal) self.assertIsInstance(ds.x, tfd.Sample) # Static properties. self.assertAllEqual( Model(e=tf.float32, scale=tf.float32, loc=tf.float32, m=tf.float32, x=tf.int32), d.dtype) batch_shape_tensor_, event_shape_tensor_ = self.evaluate( [d.batch_shape_tensor(), d.event_shape_tensor()]) expected_batch_shape = Model(e=[], scale=[], loc=[], m=[], x=[]) for (expected, actual_tensorshape, actual_shape_tensor_) in zip(expected_batch_shape, d.batch_shape, batch_shape_tensor_): self.assertAllEqual(expected, actual_tensorshape) self.assertAllEqual(expected, actual_shape_tensor_) expected_event_shape = Model(e=[2], scale=[], loc=[], m=[], x=[12]) for (expected, actual_tensorshape, actual_shape_tensor_) in zip(expected_event_shape, d.event_shape, event_shape_tensor_): self.assertAllEqual(expected, actual_tensorshape) self.assertAllEqual(expected, actual_shape_tensor_) expected_jlp = sum(d.log_prob(x) for d, x in zip(ds, xs)) actual_jlp = d.log_prob(xs) self.assertAllClose(*self.evaluate([expected_jlp, actual_jlp]), atol=0., rtol=1e-4)
def __init__(self, inputs, outputs, kernel_type, noise_level=1e-3, hyp_priors={}): # Inputs: # inputs := numpy array of inputs # outpust:= numpy array of outputs # kernel_type := string specifying the type of kernel to be used. Options are # 'RBF', 'Matern12', 'Matern32', 'Matern52' # noise_level := value for the noise variance of the output # hyp_priors := dictionary containing information about the prior distribution of the hyperparamters #------- Storing training data------------- if len(inputs.shape) == 1: self.dim_input = 1 self.Xtrain = tf.convert_to_tensor(inputs[:, None], tf.float32) else: self.dim_input = inputs.shape[1] self.Xtrain = tf.convert_to_tensor(inputs, tf.float32) self.Ytrain = tf.convert_to_tensor(outputs, tf.float32) self.n_train = len(outputs) # the number of outputs #------ Priors for the Gaussian process model----------- # Priors on the inverse lengthscale if 'beta' in hyp_priors: try: concentration = tf.convert_to_tensor( hyp_priors['beta']['concentration'], tf.float32) rate = tf.convert_to_tensor(hyp_priors['beta']['rate'], tf.float32) except Exception as e: traceback.print_exc() print( 'Could not retrieve prior distibution information for beta.' ) conc_shape = np.array(concentration.shape) rate_shape = np.array(rate.shape) correct_shape = np.array([self.dim_input]) invalid = not (np.array_equal(conc_shape, correct_shape)) or not ( np.array_equal(rate_shape, correct_shape)) if invalid: raise Exception( 'Incorrect numpy array shape for rate or concentration for beta hyperparameter.' ) else: concentration = tf.ones(self.dim_input, tf.float32) rate = tf.ones(self.dim_input, tf.float32) self.rv_beta = tfd.Independent(tfd.Gamma(concentration=concentration, rate=rate), reinterpreted_batch_ndims=1, name='rv_beta') # prior on the mean if 'loc' in hyp_priors: try: loc = tf.convert_to_tensor(hyp_priors['loc']['loc'], tf.float32) scale = tf.convert_to_tensor(hyp_priors['loc']['scale'], tf.float32) except Exception as e: traceback.print_exc() print( 'Could not retrieve prior distibution information for loc.' ) invalid = not (type(loc) == float) or not (type(scale) == float) if invalid: raise Exception( 'Incorrect type for loc or scale for loc hyperparameter. Values must be of type float.' ) else: loc = 0.0 scale = 0.5 self.rv_loc = tfd.Normal(loc=loc, scale=scale, name='rv_loc') # prior on the variance if 'varm' in hyp_priors: try: concentration = hyp_priors['varm']['concentration'] rate = hyp_priors['varm']['rate'] except Exception as e: traceback.print_exc() print( 'Could not retrieve prior distibution information for varm.' ) invalid = not (type(concentration) == float) or not (type(rate) == float) if invalid: raise Exception( 'Incorrect type for rate or concentration for varc hyperparameter. Values must be of type float.' ) else: concentration = 2.0 rate = 2.0 self.rv_varm = tfd.Gamma(concentration=concentration, rate=rate, name='rv_varm') # prior on the noise variance self.rv_noise = tfd.LogNormal(loc=-6.9, scale=1.5, name='rv_noise') #-- value for the variance of the Gaussian noise self.noise = noise_level self.jitter_level = 1e-6 # jitter level to deal with numerical instability with cholesky factorization self.kernel, self.expected_kernel = kernel_mapping[kernel_type] return
def _sample_n(self, n, seed=None): return 1. / tfd.Gamma( concentration=self.concentration, rate=self.scale).sample(n, seed=seed)
def _as_distribution(self, r): concentration = DeferredTensor(self._concentration, lambda x: tf.cast(x, r.dtype), dtype=r.dtype) return tfd.Gamma(concentration=concentration, rate=DeferredTensor(r, lambda x: tf.math.exp(-x)))
class MarkovChainBijectorTest(test_util.TestCase): # pylint: disable=g-long-lambda @parameterized.named_parameters( dict(testcase_name='deterministic_prior', prior_fn=lambda: tfd.Deterministic([-100., 0., 100.]), transition_fn=lambda _, x: tfd.Normal(loc=x, scale=1.)), dict(testcase_name='deterministic_transition', prior_fn=lambda: tfd.Normal(loc=[-100., 0., 100.], scale=1.), transition_fn=lambda _, x: tfd.Deterministic(x)), dict(testcase_name='fully_deterministic', prior_fn=lambda: tfd.Deterministic([-100., 0., 100.]), transition_fn=lambda _, x: tfd.Deterministic(x)), dict(testcase_name='mvn_diag', prior_fn=( lambda: tfd.MultivariateNormalDiag(loc=[[2.], [2.]], scale_diag=[1.])), transition_fn=lambda _, x: tfd.VectorDeterministic(x)), dict(testcase_name='docstring_dirichlet', prior_fn=lambda: tfd.JointDistributionNamedAutoBatched( {'probs': tfd.Dirichlet([1., 1.])}), transition_fn=lambda _, x: tfd.JointDistributionNamedAutoBatched( {'probs': tfd.MultivariateNormalDiag(loc=x['probs'], scale_diag=[0.1, 0.1])}, batch_ndims=ps.rank(x['probs']))), dict(testcase_name='uniform_step', prior_fn=lambda: tfd.Exponential(tf.ones([4, 1])), transition_fn=lambda _, x: tfd.Uniform(low=x, high=x + 1.)), dict(testcase_name='joint_distribution', prior_fn=lambda: tfd.JointDistributionNamedAutoBatched( batch_ndims=2, model={ 'a': tfd.Gamma(tf.zeros([5]), 1.), 'b': lambda a: ( tfb.Reshape( event_shape_in=[4, 3], event_shape_out=[2, 3, 2])( tfd.Independent( tfd.Normal( loc=tf.zeros([5, 4, 3]), scale=a[..., tf.newaxis, tf.newaxis]), reinterpreted_batch_ndims=2)))}), transition_fn=lambda _, x: tfd.JointDistributionNamedAutoBatched( batch_ndims=ps.rank_from_shape(x['a'].shape), model={'a': tfd.Normal(loc=x['a'], scale=1.), 'b': lambda a: tfd.Deterministic( x['b'] + a[..., tf.newaxis, tf.newaxis, tf.newaxis])}) ), dict(testcase_name='nested_chain', prior_fn=lambda: tfd.MarkovChain( initial_state_prior=tfb.Split(2)( tfd.MultivariateNormalDiag(0., [1., 2.])), transition_fn=lambda _, x: tfb.Split(2)( tfd.MultivariateNormalDiag(x[0], [1., 2.])), num_steps=6), transition_fn=( lambda _, x: tfd.JointDistributionSequentialAutoBatched( [ tfd.MultivariateNormalDiag(x[0], [1.]), tfd.MultivariateNormalDiag(x[1], [1.])], batch_ndims=ps.rank(x[0]))))) # pylint: enable=g-long-lambda def test_default_bijector(self, prior_fn, transition_fn): chain = tfd.MarkovChain(initial_state_prior=prior_fn(), transition_fn=transition_fn, num_steps=7) y = self.evaluate(chain.sample(seed=test_util.test_seed())) bijector = chain.experimental_default_event_space_bijector() self.assertAllEqual(chain.batch_shape_tensor(), bijector.experimental_batch_shape_tensor()) x = bijector.inverse(y) yy = bijector.forward( tf.nest.map_structure(tf.identity, x)) # Bypass bijector cache. self.assertAllCloseNested(y, yy) chain_event_ndims = tf.nest.map_structure( ps.rank_from_shape, chain.event_shape_tensor()) self.assertAllEqualNested(bijector.inverse_min_event_ndims, chain_event_ndims) ildj = bijector.inverse_log_det_jacobian( tf.nest.map_structure(tf.identity, y), # Bypass bijector cache. event_ndims=chain_event_ndims) if not bijector.is_constant_jacobian: self.assertAllEqual(ildj.shape, chain.batch_shape) fldj = bijector.forward_log_det_jacobian( tf.nest.map_structure(tf.identity, x), # Bypass bijector cache. event_ndims=bijector.inverse_event_ndims(chain_event_ndims)) self.assertAllClose(ildj, -fldj) # Verify that event shapes are passed through and flattened/unflattened # correctly. inverse_event_shapes = bijector.inverse_event_shape(chain.event_shape) x_event_shapes = tf.nest.map_structure( lambda t, nd: t.shape[ps.rank(t) - nd:], x, bijector.forward_min_event_ndims) self.assertAllEqualNested(inverse_event_shapes, x_event_shapes) forward_event_shapes = bijector.forward_event_shape(inverse_event_shapes) self.assertAllEqualNested(forward_event_shapes, chain.event_shape) # Verify that the outputs of other methods have the correct structure. inverse_event_shape_tensors = bijector.inverse_event_shape_tensor( chain.event_shape_tensor()) self.assertAllEqualNested(inverse_event_shape_tensors, x_event_shapes) forward_event_shape_tensors = bijector.forward_event_shape_tensor( inverse_event_shape_tensors) self.assertAllEqualNested(forward_event_shape_tensors, chain.event_shape_tensor())
def __init__(self, train_features, train_labels, test_features=None, test_labels=None, name='sparse_logistic_regression', pretty_name='Sparse Logistic Regression'): """Construct the sparse logistic regression model. Args: train_features: Floating-point `Tensor` with shape `[num_train_points, num_features]`. Training features. train_labels: Integer `Tensor` with shape `[num_train_points]`. Training labels. test_features: Floating-point `Tensor` with shape `[num_test_points, num_features]`. Testing features. Can be `None`, in which case test-related sample transformations are not computed. test_labels: Integer `Tensor` with shape `[num_test_points]`. Testing labels. Can be `None`, in which case test-related sample transformations are not computed. name: Python `str` name prefixed to Ops created by this class. pretty_name: A Python `str`. The pretty name of this model. Raises: ValueError: If `test_features` and `test_labels` are either not both `None` or not both specified. """ with tf.name_scope(name): num_features = int(train_features.shape[1] + 1) self._prior_dist = tfd.JointDistributionNamed( dict( unscaled_weights=tfd.Sample(tfd.Normal(0., 1.), num_features), local_scales=tfd.Sample(tfd.Gamma(0.5, 0.5), num_features), global_scale=tfd.Gamma(0.5, 0.5), )) def log_likelihood_fn(unscaled_weights, local_scales, global_scale, features, labels, reduce_sum=True): """The log_likelihood function.""" features = tf.convert_to_tensor(features, tf.float32) features = _add_bias(features) labels = tf.convert_to_tensor(labels) weights = ( unscaled_weights * local_scales * global_scale[..., tf.newaxis]) logits = tf.einsum('nd,...d->...n', features, weights) log_likelihood = tfd.Bernoulli(logits=logits).log_prob(labels) if reduce_sum: return tf.reduce_sum(log_likelihood, [-1]) else: return log_likelihood self._train_log_likelihood_fn = functools.partial( log_likelihood_fn, features=train_features, labels=train_labels) sample_transformations = { 'identity': model.Model.SampleTransformation( fn=lambda params: params, pretty_name='Identity', dtype=self._prior_dist.dtype, ) } if (test_features is not None) != (test_labels is not None): raise ValueError('`test_features` and `test_labels` must either both ' 'be `None` or both specified. Got: test_features={}, ' 'test_labels={}'.format(test_features, test_labels)) if test_features is not None and test_labels is not None: test_log_likelihood_fn = functools.partial( log_likelihood_fn, features=test_features, labels=test_labels) sample_transformations['test_nll'] = ( model.Model.SampleTransformation( fn=lambda params: test_log_likelihood_fn(**params), pretty_name='Test NLL', )) sample_transformations['per_example_test_nll'] = ( model.Model.SampleTransformation( fn=lambda params: test_log_likelihood_fn( # pylint: disable=g-long-lambda reduce_sum=False, **params), pretty_name='Per-example Test NLL', )) super(SparseLogisticRegression, self).__init__( default_event_space_bijector=dict( unscaled_weights=tfb.Identity(), local_scales=tfb.Exp(), global_scale=tfb.Exp(), ), event_shape=self._prior_dist.event_shape, dtype=self._prior_dist.dtype, name=name, pretty_name=pretty_name, sample_transformations=sample_transformations, )