allow nonlinear model callables to be time-indexed

The hmm_fillter code allows the user to pass in a transition_function that maps from time index t to a KxK transition matrix. We should add similar functionality to ekf_filter, which currently assumes the dynamics and emission functions are time-invariant (they depend only on the state and optionally the input).

As a unit test, you should reproduce the kf-linreg notebook, using a dynamics model of the form

f(z_{t-1}, u_t, t) = X_t * z_{t-1}

where X_t is the t'th feature vector (of size D), and z_t is the weight vector (of size D).
The input u_t is ignored. This is therefore a time-varying linear gaussian system.

Start with dynamax/nonlinear_gaussian_ssm/inference_ekf.py and compare its callable handling with dynamax/hidden_markov_model/inference.py. Then inspect the docs/notebooks/linear_gaussian_ssm/kf_linreg.ipynb notebook and reproduce its setup as the unit test, using the time-varying dynamics described in the issue. Done means EKF dynamics and emissions support the requested time index and the regression example is covered by a passing test.