GRU Prediction + Gradient Optimizer XLA Incompatible

[frehe]

Running the control breaks when running on the CartPoleSimulator:

• MPC Controller
• A gradient optimizer (tensorflow)
• GRU prediction model
• XLA (JIT) compilation mode

[frehe]

The problem is that there is some transpose operation when doing the GRU prediction, which XLA cannot intepret at compile time, probably because it cannot predict the dimensions after applying that op. But would need further digging to find the line of code where this is happening