This reverts commit c16f3260. There's a bug in the commit creates a scalar result with `ShapeCastOp`. Reverting till that fix is done.