[XLA] Do not recompute flattened sets inside layout assignment.
Cache the flattened sets instead of recomputing them. This matters for large graphs, since we may request the flattened set thousands of times on the same instruction, and it may be fairly expensive to construct for large tuples. PiperOrigin-RevId: 187046642
Loading
Please sign in to comment