tree d943ba0de70f48b5fe062bfb46f28d45f9faef24
parent 2780589f2cfa635b8d5713373835621b17e11e0f
author Eugene Burmako <burmako@google.com> 1658716220 -0700
committer TensorFlower Gardener <gardener@tensorflow.org> 1658716443 -0700

Extend the canonicalization rules for dynamic_iota.

Recognize a dynamic_iota whose argument is a constant,
and turn it into an iota.

I hoped we can get this done via shape inference, using the newly introduced
refineReturnTypes from MLIR upstream, but that doesn't work because we cannot
define inferReturnTypes for dynamic_iota, and it's still used in a number of
places.

PiperOrigin-RevId: 462987927