Error in documentation of experimental.sequential.extended_kalman_filter

Hi all,

there are two small issues with the documentation of experimental.sequential.extended_kalman_filter.

1. The code in the example does not run, giving

InvalidArgumentError: cannot compute Sub as input #1(zero-based) was expected to be a double tensor but is a float tensor [Op:Sub] name.

Changing

x = [np.zeros((2,), dtype=np.float32)]
to
x = [tf.zeros((2,), dtype=tf.float32)]
solves this.

2. The documentation states

   observation_jacobian_fn: a Python callable that accepts a (batched) vector
   of length state_size and returns a (batched) matrix of size
   [state_size, event_size], representing the Jacobian of observation_fn.

but the correct matrix size to be returned should be [event_size, state_size].

In the example, making observation_jacobian_fn return a Tensor of size [state_size, event_size] will cause the Kalman Filter to fail, while [event_size, state_size] works.

See also this colab https://colab.research.google.com/drive/1vUNMHQ1Fc3CwdpOMnKXqzqmQ2zM4A8Pp?usp=sharing

Thank you for the great package!

Open the documentation for experimental.sequential.extended_kalman_filter and run its example, using the linked Colab as a reproduction reference. Update the example's array construction and correct the documented observation_jacobian_fn matrix shape, then rerun the example to confirm it works.