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
Loading
Please sign in to comment