def _init_distribution(conditions, **kwargs):
loc, concentration = conditions["loc"], conditions["concentration"]
return tfd.InverseGaussian(loc=loc,
concentration=concentration,
**kwargs)
import tensorflow_probability as tfp
from functools import *
from tensorflow_probability import bijectors as tfb
from tensorflow_probability import distributions as tfd
from tensorflow_probability.python.internal import dtype_util
from tensorflow_probability.python.internal import prefer_static as ps
from tensorflow_probability.python.internal import tensorshape_util
tf.enable_v2_behavior()
X = tf.constant(1.0)
scale1 = 1.2
alpha = 12.3
beta = 0.2
mean = 0.0
x = tf.Variable(2.2)
pd = tfd.JointDistributionSequential([
tfd.Independent(tfd.InverseGaussian(scale1, (alpha**2 - beta**2)**.5),
reinterpreted_batch_ndims=0),
lambda mix: tfd.Normal(loc=mean + beta * mix, scale=mix)
])
def _make_val_and_grad_fn(value_fn):
@functools.wraps(value_fn)
def val_and_grad(x):
return tfp.math.value_and_gradient(value_fn, x)
return val_and_grad
def nig(mean, scale1, alpha, beta):
pd = tfd.JointDistributionSequential([
def nig(mean, scale1, alpha, beta):
pd = tfd.JointDistributionSequential([
tfd.Independent(tfd.InverseGaussian(scale1, (alpha**2 - beta**2)**.5)),
lambda mix: tfd.Normal(loc=mean + beta * mix, scale=mix)
])
return pd.log_prob([pd.sample()[0], x])
def _base_dist(self, mu: TensorLike, lam: TensorLike, *args, **kwargs):
return tfd.InverseGaussian(loc=mu, concentration=lam, *args, **kwargs)