Calling an external function - the posterior distribution of a gaussian process fit


Basically the question is: how do I use a function which does not act on pytensors, but instead on tensorflow tensors?

I am using a gaussian process to fit a function with GPFlow which gives me attenuation as a function of absorption. Absorption is due to a few things which make up the material, so I place priors on the constituents and do a linear combination taking into account the varying absorption at different wavelengths.

# some numpy array with the same shape as data
wavelength_dependent_absorption_coefficient = [...]
with pm.Model() as model:
  constituent = pm.Normal('constituent', 1, 0.1)
  absorption = wavelength_dependent_absorption_coefficient * constituent

  y = pm.Normal('y', mu=gaussian_process(absorption), sigma=0.1, observed=data)

Unfortunately this doesn’t work, giving an error like: ValueError: setting an array element with a sequence. This issue is discussed inside the ‘using a “black box” likelihood’ example. However, ideally, I’d like to find a simpler solution. I’m very willing to implement the gaussian process in pymc but at the moment would have the same problem with calling an external function.

You have to wrap your function in an Op, which the blackbox likelihood notebook demonstrates, there is no other way around.

Here are some examples with wrapping JAX functions that may feel more succinct: How to use JAX ODEs and Neural Networks in PyMC - PyMC Labs