"log_likelihood" not found in InferenceData

Questions
4

I am using ArviZ version 0.11.4, PyMC3 version 3.11.4.

Here is a minimal example in which the log_likelihood is not saved in the InferenceData (i.e., not in trace).

import arviz as az
import matplotlib.pyplot as plt
import numpy as np
import pymc3 as pm
import theano.tensor as tt
from matplotlib import pyplot

NUM_CORES = 4
NUM_DRAWS = 25000

class BlackBoxLikelihood(tt.Op):
    '''
    Specify what type of object will be passed and returned to the Op when it is
    called. In our case we will be passing it a vector of values (the parameters
    that define our model) and returning a single "scalar" value (the
    log-likelihood)
    '''
    itypes = [tt.dvector] # Expects a vector of parameter values when called
    otypes = [tt.dscalar] # Outputs a single scalar value (the log likelihood)

    def __init__(self, model, observed, x, loglik = None):
        '''
        Initialise the Op with various things that our log-likelihood function
        requires. Below are the things that are needed in this particular
        example.

        Parameters
        ----------
        model : Callable
            An arbitrary "black box" function that takes two arguments: the
            model parameters ("params") and the forcing data ("x")
        observed : numpy.ndarray
            The "observed" data that our log-likelihood function takes in
        x:
            The forcing data (input drivers) that our model requires
        loglik : Callable or None
            The log-likelihood (or whatever) function we've defined
        '''
        self.model = model
        self.observed = observed
        self.x = x
        # Optionally overwrite the Gaussian log-likelihood function
        if loglik is not None:
            self.loglik = loglik

    def loglik(self, params, x, observed):
        predicted = self.model(params, x)
        sigma = params[-1]
        return -0.5 * np.log(2 * np.pi * sigma**2) - (0.5 / sigma**2) *\
            np.nansum((predicted - observed)**2)

    def perform(self, node, inputs, outputs):
        # The method that is used when calling the Op
        (params,) = inputs
        logl = self.loglik(params, self.x, self.observed)
        outputs[0][0] = np.array(logl) # Output the log-likelihood


def main():
    def model(params, x, n = 1000):
        'Our black-box model; really, a simulator of Gaussian random vars'
        return np.random.normal(*params, n)

    # Observed data
    data = np.random.normal(loc = 0, scale = 1, size = 1000)
    loglik = BlackBoxLikelihood(model, data, None)

    basic_model = pm.Model()
    with basic_model:
        # (Stochstic) Priors for unknown model parameters
        a = pm.Normal('a', mu = 0, sd = 5)
        b = pm.HalfNormal('b', sd = 1)

        # Convert model parameters to a tensor vector
        params = tt.as_tensor_variable([a, b])
        pm.Potential('likelihood', loglik(params))

        trace = pm.sample(
            draws = NUM_DRAWS, step = pm.DEMetropolisZ(vars = [a, b]),
            cores = NUM_CORES, chains = NUM_CORES,
            return_inferencedata = True, start = {'a': 1, 'b': 10},
            idata_kwargs = {'log_likelihood': True})

Output trace looks like the following, even if idata_kwargs is omitted:

>>> trace
Inference data with groups:
	> posterior
	> sample_stats