tree 65e9efb26ec55cdd92f57f0ae811bb213b53bca8
parent 89133a6a74da4b976a585b66a09a6601fce3b217
author Brian Patton <bjp@google.com> 1618410486 -0700
committer TensorFlower Gardener <gardener@tensorflow.org> 1618410728 -0700

Look through get-tuple-element(tuple(..), i) instructions when resolving constants in xla_builder.

Update tf2xla_util to rewrite the while body so that replaced loop-invariant _Arg nodes feed straight to the corresponding _Retval. This ensures IsLoopInvariant will continue to return true even after the rewrite.

These changes allow a new test to pass in def_function_xla_jit_test.

PiperOrigin-RevId: 368425851
Change-Id: I33ed4149b4c1069e13eab469f9bc853c00e595ce