Commit a6ed38fe authored by Justin Lebar's avatar Justin Lebar Committed by TensorFlower Gardener
Browse files

[XLA] Set layout of GTE instructions inside fusion nodes.

Other than the root and parameters of a fusion computation, most other
instructions in a fusion computation don't have a layout.  GTEs are an
exception; they should inherit their layout from their operand, which
must be another GTE or a parameter.

Previously LayoutAssignment left GTEs alone, assuming they came in with
the right layout.  But this isn't correct, and in fact LayoutAssignment
cleared the layouts of every non-fused instruction before assigning them
for exactly this reason.  If we'd done the same to fused instructions,
it would have caught this bug, so we make that change here as well.  (We
simplify this loop by removing the check for kOutfeed -- outfeeds do not
produce a result, so there's no shape to keep.)

PiperOrigin-RevId: 183595627
parent 74b5a4cf
Loading
Loading
Loading
Loading
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please to comment