[XLA] Generalize pooling to support N-dimensional inputs.
Extends the usage of windowed pipeline emitter in pooling to support arbitrary input dimensions > 3. Users should note that memory usage compounds quickly with very-high dimensional tensors. Adds an HLO evaluator test for the R6 pooling case. PiperOrigin-RevId: 167040226
Loading
Please sign in to comment