Elliptical Slice Sampling with Latent Covariance Matrix

I am trying to perform a version of elliptical slice sampling in a model where the data-points, and hence the covariance matrix, are modelled as latent variables drawn from a uniform distribution. Here is the relevant code:

N = 10 # Set number of samples
gp_ls = 1 # Lengthscale for the multivariate normal prior over phis

with pm.Model() as pref_model:
    
    # Prior distribution over policies
    pols = pm.Uniform('pols', lower=0, upper=1, shape=(N, 2))
    
    pols_unrolled = pols.T.reshape((2 * N, 1)) # Ordered as left pols, then right pols
                
    # 'Prior' distribution over phis (given policy values) (GP version)
    cov_func = pm.gp.cov.ExpQuad(input_dim=1, ls=gp_ls) # Each policy is 1-dimensional
    gp = pm.gp.Latent(cov_func=cov_func)
    phis = gp.prior('phis', X=pols_unrolled, shape=(2 * N, 1), reparameterize = False)
    
    
    K = cov_func(pols_unrolled)
    step1 = pm.NUTS(vars=[pols])
    step2 = pm.EllipticalSlice(vars=[phis], prior_cov=K)
    trace = pm.sample(cores=4, steps=[step1, step2]) 

However, this results in the following error:

TypeError: For compute_test_value, one input test value does not have the requested type.

The error when converting the test value to that variable type:
Non-unit value on shape on a broadcastable dimension.
(20, 20)
(False, True)

I believe this has something to do with the matrix K being recalculated at every sample within the model, rather than it being a fixed matrix based on observed variables. Are there any ways of resolving this? If I use NUTS for all of the variables, it works fine.

Elliptical slice sampling might be moved into a more experimental section of the repo. From taking a look at the code, it looks like the covariance matrix used in it currently has to not be a random variable, but fixed. This does limit its usefulness. Iā€™m not sure what effort would be required to allow it to be able to use a random variable, as is your case. PR welcome! I do think using NUTS is a good approach here. My understanding of the research is that HMC/NUTS on a reparameterized version of the covariance is currently preferred over elliptical slice sampling for this scenario, but I could be wrong.

1 Like