How best to optimize a function used in DensityDist?

I have been optimizing my function “VectorizedVcdfe”, and I am still trying to optimize it. This function is responsible for 99% of the slowness of another function “customFunc”. This “customFunc” is used as an argument of “pymc3.DensityDist”.

Please help me optimize “VectorizedVcdfe”.

Here is the function to optimize:

def VectorizedVcdfe(self, x, dataVector, recip_h_times_lambda_vector):
    n = len(dataVector)

    xVector = pymc3.math.stack([x for i in range(n)])
    differenceVector = xVector - dataVector

    stackedDiffVecAndRecipVec = pymc3.math.stack(differenceVector, recip_h_times_lambda_vector)

    erfcTerm  = 1. - pymc3.math.erf(self.neg_sqrt1_2 * pymc3.math.prod(stackedDiffVecAndRecipVec, axis=0))

    # Calc F_Hat
    F_Hat = (1. / float(n)) * pymc3.math.sum(0.5 * erfcTerm)

    # Return F_Hat
    return(F_Hat)

Arguments/variables can be described as follows…

“x” is a TensorVariable.

“dataVector” is a 1Xn numpy matrix.

“recip_h_times_lambda_vector” is also a 1Xn numpy matrix.

“neg_sqrt1_2” is a scalar constant.

Here is how “customFunc” is used:

with pymc3.Model() as model:
    # Create likelihood
    like = pymc3.DensityDist('X', customFunc, shape=2)

    # Make samples
    step = pymc3.NUTS()
    trace = pymc3.sample(2000, tune=1000, init=None, step=step, cores=2)

The general suggestion is to move all the numpy operation outside and avoid using for loop, for example, likely you can remove the below two lines

    n = len(dataVector)

    xVector = pymc3.math.stack([x for i in range(n)])

as xVector - dataVector should broadcast correctly.

Thank you, that reduced the time by a factor of 9

1 Like