Examples broken in rewrite branch

[adam-hartshorne]

I am sure you are probably aware, but the basic examples like ode_demo and latent_ode are broken in the rewrite branch. The size of the tensor returned from odeint is different from the input, and thus within the training loop pred_y shape doesn't equal the batch_y shape and causes an error when calculating the loss function.

I am only attempting to use this branch as I would like to test a Neural ODE with some code I already have written with TF 2.3.

[titu1994]

Please use just the master branch. The rewrite branch does not work, yet.