Commit e2e67c52 authored by Sanjoy Das's avatar Sanjoy Das Committed by TensorFlower Gardener
Browse files

Fix some edge cases around scalar indices in the gather expander

I discovered these when changing the tf2xla bridge to directly emit gather
operations.

 - DeScalarizeGatherIndices was assuming that gather_indices must be of at least
   rank 1.  Fix this to be more general.

 - We were passing in the wrong version of gather indices to
   ExpandFirstDimIntoNDims.  We don't strictly need to pass in
   transposed_gather_indices (since if transposed_gather_indices is rank 1 then
   the transpose has to be an identity transpose), passing in
   descalarized_gather_indices would also have been fine, but
   transposed_gather_indices seems more uniform.

 - ExpandGatherDimsInAccumulator was assuming that gather_indices must be of at
   least rank 1 (by calling CollapseFirstNDims).  Fix this to be more general.

 - We were trying to go through with emitting zero sized gather operations.  I
   don't think it is worth dealing with all of the edge cases this would expose
   so now we just punt to ZeroSizedHloElimination.

PiperOrigin-RevId: 189696444
parent 2311e9ce
Loading
Loading
Loading
Loading
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please to comment