Use casting and tf.gather to avoid end-to-end XLA failure

Merged Chris Jewell requested to merge github/fork/csuter/stochastic-da-dev into stochastic-da-dev

Created by: csuter

Merge request reports