Commit bbe996e2 authored by Cong Liu's avatar Cong Liu Committed by TensorFlower Gardener
Browse files

[TF2XLA] Use xla::GetDimensionSize in reduce_mean. 

In the future when XLA has dynamic shapes, the MeanOp needs to divide the real
dimension size.

[XLA] Add a HloGetDimensionSizeRewritter pass which converts a
GetDimensionSize op into a Constant.
The lowering of GetDimensionSize is not yet implemented on any platforms. So
add a pass that replaces it with Constant.

PiperOrigin-RevId: 222144256
parent 1bd82e09
Show whitespace changes
Inline Side-by-side
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Please to comment