tf-probability error using Wishart prior

[doctorwes]

I am having trouble estimating a covariance matrix using a Wishart prior. This may be related to a previously reported issue in tensorflow-probability: https://github.com/GPflow/GPflow/issues/553

Error message:

InvalidArgumentError: Cholesky decomposition was not successful. The input might not be valid. [[{{node cov_194/log_prob/Cholesky}}]]

Code:

def flat_model(mean=mean0, cov=cov0):
    meanrets = inf.Normal(mean0, scale=0.01, name='meanrets')
    cov = inf.Wishart(df=n, scale=cov0, name='cov')
    with inf.datamodel():
        x = inf.MultivariateNormalFullCovariance(loc=meanrets, covariance_matrix=cov, name='x')
 
@inf.probmodel
def flat_qmodel():
    q_means_loc = inf.Parameter(np.zeros([n]), name='q_means_loc')
    q_means_scale = tf.math.softplus(inf.Parameter(np.ones([n]), name='q_means_scale'))
    qmeans = inf.Normal(q_means_loc, q_means_scale, name='meanrets')  
    q_cov_scale = inf.Parameter(np.diag(n*[1.]), name='q_cov_scale')
    qcov = inf.Wishart(df=n, scale=q_cov_scale, name='cov')```

[andresmasegosa]

There are several issues with the use of Wishart distributions. The first one seems to be the use of float32 precision, then Cholesky decomposition may easily fails,

UCL-SML/Doubly-Stochastic-DGP#9

This code does not raise this error. Note the use of inf.set_floatx('float64')' to modify the precision. Note also how cov0 is defined as well as the qmodel to guarantee a proper cholesky decomposition of the data.

`
import inferpy as inf
import tensorflow as tf
import tensorflow_probability as tfp
import numpy as np

inf.set_floatx('float64')

n=10
mean0=np.ones([n])
cov0=np.diag(n*[1.])

@inf.probmodel
def flat_model(mean=mean0, cov=cov0):
meanrets = inf.Normal(mean, scale=0.01, name='meanrets')
cov = inf.Wishart(df=n, scale=cov, name='cov')
with inf.datamodel():
x = inf.MultivariateNormalFullCovariance(loc=meanrets, covariance_matrix=cov, validate_args = True, name='x')

@inf.probmodel
def flat_qmodel():
q_means_loc = inf.Parameter(np.zeros([n]), name='q_means_loc')
q_means_scale = tf.math.softplus(inf.Parameter(np.ones([n]), name='q_means_scale'))
qmeans = inf.Normal(q_means_loc, q_means_scale, name='meanrets')
q_cov_scale = tf.matrix_diag(tf.math.softplus(inf.Parameter(np.zeros([n])),name='q_cov_scale'))
qcov = inf.Wishart(df=n, scale=q_cov_scale, name='cov')

m = flat_model()

data = m.prior().sample(100)["x"]
optimizer = tf.train.AdamOptimizer(0.01)
VI = inf.inference.VI(flat_qmodel(), optimizer=optimizer, epochs=5000)

m.fit({"x": data}, VI)

m.posterior("meanrets").parameters()
m.posterior("cov").parameters()
`

Even though it seems the VI algorithm does not converge. This notebook analyze the problem of covariance estimation using a Bayesian approach in detail. By it is really involve and requires of the use of bijectors.

https://colab.research.google.com/github/tensorflow/probability/blob/master/tensorflow_probability/examples/jupyter_notebooks/TensorFlow_Probability_Case_Study_Covariance_Estimation.ipynb

If we find a simpler way to achieve it, we will include it in the documentation.