#332 Closed form gradients for Kalman filter

[JeanVanDyk]

Here you'll find, in the notebook section, the notebook I've used to compare execution time between using autodiff and the analytic gradients.

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[ricardoV94]

Suggestion, run custom gradient vs default autodiff in pytensor and time it. Print the compiled graph (.dprint() on the function) to see if what pytensor autodiff is doing isn't already similar.

I see the original paper compares with PyTorch, which IMO isn't very clever AD, specially in terms of memory optimization. I'm not sure their 38x speedup / memory improvement also holds against pytensor (or jax).

[JeanVanDyk]

Thanks @ricardoV94!

I’ve tried looking into it, but it quickly becomes messy… Do you have a particular method for comparing the dprint output? I saw that you can name operations to make things clearer, but that doesn’t seem very efficient given the hundreds of lines I get.

I also tried timing it, and the gradient with the closed-form expression actually performs worse. I’m inclined to think that autodiff is still being used under the hood, especially since the runtime scales in pretty much the same way for both forms, depending on the number of states.

[ricardoV94]

Can you share the timing code, that's easier to give feedback over.

Re the dprint that was just curiosity just paste it after you compile (yeah it will be long)

[JeanVanDyk]

You can find it at the end of the notebook, here is the benchmark function :

def benchmark_kalman_gradients(loss, obs_data, a0, P0, T, Z, R, H, Q):
    results = defaultdict(dict)
    exec_time = 0

    grad_list = pt.grad(loss, [data_sym, a0_sym, P0_sym, T_sym, Z_sym, H_sym, Q_sym])
    f_grad = pytensor.function(
        inputs=[data_sym, a0_sym, P0_sym, T_sym, Z_sym, H_sym, Q_sym],
        outputs=grad_list,
    )

    for _ in range(20):
    
        # --- exécution ---
        t0 = perf_counter()
        _ = f_grad(
            obs_data[:, np.newaxis],
            a0,
            P0,
            T,
            Z,
            H,
            R @ Q @ R.T,
        )
        t1 = perf_counter()
        exec_time += (t1 - t0)/20
    
    
    results["exec_time"] = exec_time

    return results
    ```