Since you have a stochastic function, I would try to rewrite it as a custom distribution. It would need some thought of how to express the output of the said stochastic function as samples from some distribution. Try to write down what are the input and the output of the original function, and then rewrite it into something like p(output | input)