From ae41dbfad55a26426c6c30ea94c341126211bb6c Mon Sep 17 00:00:00 2001 From: GianluigiSilvestri Date: Thu, 3 Jun 2021 11:51:13 +0200 Subject: [PATCH 1/2] added first version of gaussian binary tree --- .../targets/gaussian_binary_tree.py | 103 ++++++++++++++++++ .../targets/gaussian_binary_tree_test.py | 0 2 files changed, 103 insertions(+) create mode 100644 spinoffs/inference_gym/inference_gym/targets/gaussian_binary_tree.py create mode 100644 spinoffs/inference_gym/inference_gym/targets/gaussian_binary_tree_test.py diff --git a/spinoffs/inference_gym/inference_gym/targets/gaussian_binary_tree.py b/spinoffs/inference_gym/inference_gym/targets/gaussian_binary_tree.py new file mode 100644 index 0000000000..26f5139e99 --- /dev/null +++ b/spinoffs/inference_gym/inference_gym/targets/gaussian_binary_tree.py @@ -0,0 +1,103 @@ +import functools + +import tensorflow.compat.v2 as tf +from inference_gym.targets import bayesian_model +from inference_gym.targets import model + +import tensorflow_probability as tfp + +tfb = tfp.bijectors +tfd = tfp.distributions + +Root = tfd.JointDistributionCoroutine.Root + +# coupling link could be e.g. tf.nn.tanh +def gaussian_binary_tree_prior_fn(num_layers, initial_scale, nodes_scale, + coupling_link=None): + initial_loc = 0. + nodes = [] + # in the "root" layer (or inverse root, as it is a reversed tree) we have + # 2**num_layers nodes (with depth 2 --> 4 nodes, depth 4 --> 16 nodes) + for i in range(2 ** num_layers): + node = yield Root(tfd.Normal(initial_loc, initial_scale)) + nodes.append(node) + # for the remaining layers, we then sample the respective nodes values + # applying the link function + # we do not do this for the final node, as it is supposed to be observed + for l in range(num_layers, 1, -1): + next_layer_nodes = [] + for i in range(0, (2 ** l), 2): + if coupling_link: + node = yield tfd.Independent( + tfd.Normal(coupling_link(nodes[i]) - coupling_link(nodes[i + 1]), + nodes_scale), 0) + else: + node = yield tfd.Independent( + tfd.Normal(nodes[i] - nodes[i + 1], + nodes_scale), 0) + next_layer_nodes.append(node) + nodes = next_layer_nodes + + +def gaussian_binary_tree_log_likelihood_fn(values, observed_last_node, + nodes_scale, coupling_link=None): + left_node, right_node = values[-2], values[-1] + if coupling_link: + lps = tfd.Normal(loc=coupling_link(left_node) - coupling_link(right_node), + scale=nodes_scale).log_prob(observed_last_node) + else: + lps = tfd.Normal(loc=left_node - right_node, + scale=nodes_scale).log_prob(observed_last_node) + return lps + + +class GaussianBinaryTree(bayesian_model.BayesianModel): + def __init__(self, + num_layers, + observed_last_node, + initial_scale, + nodes_scale, + coupling_link=None, + name='gaussian_binary_tree', + pretty_name='Gaussian Binary Tree'): + """Construct the Gaussian Binary Tree model.""" + with tf.name_scope(name): + self._prior_dist = tfd.JointDistributionCoroutine(functools.partial( + gaussian_binary_tree_prior_fn, + num_layers=num_layers, + initial_scale=initial_scale, + nodes_scale=nodes_scale, + coupling_link=coupling_link + )) + + self._log_likelihood_fn = functools.partial( + gaussian_binary_tree_log_likelihood_fn, + observed_last_node=observed_last_node, + nodes_scale=nodes_scale, + coupling_link=coupling_link + ) + + # todo: what should I use here? + sample_transformations = { + 'identity': + model.Model.SampleTransformation( + fn=lambda params: params, + pretty_name='Identity', + dtype=self._prior_dist.dtype, + ) + } + + super(GaussianBinaryTree, self).__init__( + default_event_space_bijector=tfb.Identity(), # todo: what should I use here? + event_shape=self._prior_dist.event_shape, + dtype=self._prior_dist.dtype, + name=name, + pretty_name=pretty_name, + sample_transformations=sample_transformations + ) + + def _prior_distribution(self): + return self._prior_dist + + def _log_likelihood(self, value): + return self._log_likelihood_fn(value) \ No newline at end of file diff --git a/spinoffs/inference_gym/inference_gym/targets/gaussian_binary_tree_test.py b/spinoffs/inference_gym/inference_gym/targets/gaussian_binary_tree_test.py new file mode 100644 index 0000000000..e69de29bb2 From 7172acff737d03ffc3b2655346bdc7be9667b40b Mon Sep 17 00:00:00 2001 From: GianluigiSilvestri Date: Mon, 21 Jun 2021 11:10:04 +0200 Subject: [PATCH 2/2] create nodes at the same layer with sample instead of loop --- .../targets/gaussian_binary_tree.py | 38 +++++++++---------- 1 file changed, 19 insertions(+), 19 deletions(-) diff --git a/spinoffs/inference_gym/inference_gym/targets/gaussian_binary_tree.py b/spinoffs/inference_gym/inference_gym/targets/gaussian_binary_tree.py index 26f5139e99..7a836020fc 100644 --- a/spinoffs/inference_gym/inference_gym/targets/gaussian_binary_tree.py +++ b/spinoffs/inference_gym/inference_gym/targets/gaussian_binary_tree.py @@ -11,32 +11,31 @@ Root = tfd.JointDistributionCoroutine.Root + # coupling link could be e.g. tf.nn.tanh def gaussian_binary_tree_prior_fn(num_layers, initial_scale, nodes_scale, coupling_link=None): initial_loc = 0. - nodes = [] # in the "root" layer (or inverse root, as it is a reversed tree) we have - # 2**num_layers nodes (with depth 2 --> 4 nodes, depth 4 --> 16 nodes) - for i in range(2 ** num_layers): - node = yield Root(tfd.Normal(initial_loc, initial_scale)) - nodes.append(node) + # 2**(num_layers-1) nodes (with depth 2 --> 2 nodes, depth 4 --> 8 nodes) + + nodes = yield Root( + tfd.Sample(tfd.Normal(initial_loc, initial_scale), 2 ** (num_layers - 1), + name=f'layer_{num_layers - 1}')) # for the remaining layers, we then sample the respective nodes values # applying the link function # we do not do this for the final node, as it is supposed to be observed for l in range(num_layers, 1, -1): - next_layer_nodes = [] - for i in range(0, (2 ** l), 2): - if coupling_link: - node = yield tfd.Independent( - tfd.Normal(coupling_link(nodes[i]) - coupling_link(nodes[i + 1]), - nodes_scale), 0) - else: - node = yield tfd.Independent( - tfd.Normal(nodes[i] - nodes[i + 1], - nodes_scale), 0) - next_layer_nodes.append(node) - nodes = next_layer_nodes + if coupling_link: + nodes = yield tfd.Independent(tfd.Normal(tf.stack( + [coupling_link(nodes[..., i]) - coupling_link(nodes[..., i + 1]) for i + in range(0, 2 ** l, 2)], + -1), + nodes_scale), 1, name=f'layer_{l - 1}') + else: + nodes = yield tfd.Independent(tfd.Normal(tf.stack( + [nodes[..., i] - nodes[..., i + 1] for i in range(0, 2 ** l, 2)], -1), + nodes_scale), 1, name=f'layer_{l - 1}') def gaussian_binary_tree_log_likelihood_fn(values, observed_last_node, @@ -88,7 +87,8 @@ def __init__(self, } super(GaussianBinaryTree, self).__init__( - default_event_space_bijector=tfb.Identity(), # todo: what should I use here? + default_event_space_bijector=tfb.Identity(), + # todo: what should I use here? event_shape=self._prior_dist.event_shape, dtype=self._prior_dist.dtype, name=name, @@ -100,4 +100,4 @@ def _prior_distribution(self): return self._prior_dist def _log_likelihood(self, value): - return self._log_likelihood_fn(value) \ No newline at end of file + return self._log_likelihood_fn(value)