-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathmixture_of_gaussian_bijector.py
115 lines (93 loc) · 3.45 KB
/
mixture_of_gaussian_bijector.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
import tensorflow as tf
import tensorflow_probability as tfp
from tensorflow_probability.python.internal import cache_util
tfb = tfp.bijectors
tfd = tfp.distributions
tfe = tfp.experimental
class MixtureOfGaussians(tfb.Bijector):
_cache = cache_util.BijectorCacheWithGreedyAttrs(
forward_name='_augmented_forward', inverse_name='_augmented_inverse')
def __init__(self, dist, validate_args=False,
name='mixture_of_gaussians'):
# todo: assert dist is mixture of gaussians
self.dist = dist
super(MixtureOfGaussians, self).__init__(
validate_args=validate_args,
forward_min_event_ndims=0,
name=name)
def forward_pass(self, x):
x = self.dist.cdf(x)
return x
def _augmented_forward(self, x):
fldj = self.dist.log_prob(x)
return self.forward_pass(x), {'ildj': -fldj, 'fldj': fldj}
def _augmented_inverse(self, y):
bij = tfe.bijectors.ScalarFunctionWithInferredInverse(
lambda e: self.forward_pass(e), max_iterations=100, root_search_fn=tfp.math.find_root_chandrupatla)
ildj = bij.inverse_log_det_jacobian(y)
y = bij.inverse(y)
return y, {'ildj': ildj, 'fldj': -ildj}
def _forward(self, x):
y, _ = self._augmented_forward(x)
return y
def _inverse(self, y):
x, _ = self._augmented_inverse(y)
return x
def _forward_log_det_jacobian(self, x):
cached = self._cache.forward_attributes(x)
# If LDJ isn't in the cache, call forward once.
if 'fldj' not in cached:
_, attrs = self._augmented_forward(x)
cached.update(attrs)
return cached['fldj']
def _inverse_log_det_jacobian(self, y):
cached = self._cache.inverse_attributes(y)
# If LDJ isn't in the cache, call inverse once.
if 'ildj' not in cached:
_, attrs = self._augmented_inverse(y)
cached.update(attrs)
return cached['ildj']
class InverseMixtureOfGaussians(tfb.Bijector):
_cache = cache_util.BijectorCacheWithGreedyAttrs(
forward_name='_augmented_forward', inverse_name='_augmented_inverse')
def __init__(self, dist, validate_args=False,
name='mixture_of_gaussians'):
# todo: assert dist is mixture of gaussians
self.dist = dist
super(InverseMixtureOfGaussians, self).__init__(
validate_args=validate_args,
forward_min_event_ndims=0,
name=name)
def forward_pass(self, x):
x = self.dist.cdf(x)
return x
def _augmented_inverse(self, y):
ildj = self.dist.log_prob(y)
return self.forward_pass(y), {'ildj': ildj, 'fldj': -ildj}
def _augmented_forward(self, x):
bij = tfe.bijectors.ScalarFunctionWithInferredInverse(
lambda e: self.forward_pass(e), max_iterations=100,
root_search_fn=tfp.math.find_root_chandrupatla)
fldj = bij.inverse_log_det_jacobian(x)
x = bij.inverse(x)
return x, {'fldj': fldj, 'ildj': -fldj}
def _forward(self, x):
y, _ = self._augmented_forward(x)
return y
def _inverse(self, y):
x, _ = self._augmented_inverse(y)
return x
def _forward_log_det_jacobian(self, x):
cached = self._cache.forward_attributes(x)
# If LDJ isn't in the cache, call forward once.
if 'fldj' not in cached:
_, attrs = self._augmented_forward(x)
cached.update(attrs)
return cached['fldj']
def _inverse_log_det_jacobian(self, y):
cached = self._cache.inverse_attributes(y)
# If LDJ isn't in the cache, call inverse once.
if 'ildj' not in cached:
_, attrs = self._augmented_inverse(y)
cached.update(attrs)
return cached['ildj']