Python overheads of jitted functions?

[Qiustander]

Hi all I implemented the Kalman filter which is based on TFP JAX substrates. Then I return all results wrapped with a namedtuple (for example, predicted_mean, predicted_covariance). Then I wrap the whole function like this

def run_kf(lgssm):
    def inner_run(observation):
        g = kalman_filter(lgssm, observation)
        return g

    return inner_run
kf_fn = jit(run_kf(ssm_model))

The return values of kalman_filter are optional:

        if log_likelihood:
            return tf.reduce_sum(log_marginal_likelihood, axis=0)
        return KFResults(
            filtered_means, filtered_covs, predicted_means, predicted_covs, log_marginal_likelihood
        )

Then I find a very confusing phenomenon about the runtime. I tested with some cases

1D Kalman filter with 200 time steps (11 matrix multiplication): 0.002s if I only return the last likelihood which is a scalar with log_likelihood=True
1D Kalman filter with 200 time steps (11 matrix multiplication): 0.02s if log_likelihood=False
3D Kalman filter with 200 time steps (33 matrix multiplication): 0.045s if log_likelihood=True
3D Kalman filter with 200 time steps (33 matrix multiplication): 0.069s if log_likelihood=False

So it seems that there exists overhead that in caching? could we reduce this time? Thanks!

[jakevdp]

Hi - thanks for the question. I think in order to answer this we'll need more information: there's too much guesswork in trying to understand what's going on given the code you've shown.

Can you edit your question to add a minimal reproducible example?