[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
Loading
Please sign in to comment