Add slight speed improvement by avoiding tiling and reshapes.

[copybara-service[bot]]

Add slight speed improvement by avoiding tiling and reshapes.

We implicitly tile with the BE's first einsum op. Instead of reshape, we carry around the additional ensemble size axis, and where all layers operate with multiple batch dimensions.