Custom Covariance Function for Gaussian Process

Questions version agnostic
4 
    # Generate mask using the second column (rewarded_array)
    mask = pt.zeros_like(K_se)
    for i in range(X.shape[0]):
        if X[i, 1] == 1:
            mask = pt.set_subtensor(mask[i, :], 1)
            mask = pt.set_subtensor(mask[:, i], 1)

    return K_se * mask

This bit looks ripe for optimization. Generally, for loop are problematic in PyMC because they cannot be compiled with pytensor. So you’ll get most of your code executing at machine speed and then a chunk executing at python speed. We might be able to use something like:

bool_vector = pt.eq(X[:,1],1)
mask = pt.switch(bool_vector, [IF RESULT], [ELSE RESULT])

Could you say a bit about what you are trying to accomplish with:

        if X[i, 1] == 1:
            mask = pt.set_subtensor(mask[i, :], 1)
            mask = pt.set_subtensor(mask[:, i], 1)

It’s not jumping out to me why you might immediately write over the first mask with the second.

4 Likes