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)