tree 14c5dafa463a8df99666d5f4dc5b6ee91243ff67
parent 5097e6b36350f16032db950792c54f50a5385ec9
author Smit Hinsu <hinsu@google.com> 1615583891 -0800
committer TensorFlower Gardener <gardener@tensorflow.org> 1615584308 -0800

Generalize tf-rewrite-tpu-embedding-ops pass to handle embedding ops within regions

Currently, the pass assumes embedding ops to be directly in a function and was working only because functional control flow to regions pass introduces call functions so the embedding ops ended up being in a separate function.

The above won't be true for this pass's usage to rewrite these ops in MLIR module before converting it to Graph for compilation using the old bridge second phase.

PiperOrigin-RevId: 362586719
Change-Id: Ieaabcc0db8a44f6c721d1842da0720e8d9389797