Enable reshape of _ScopedAllocatorConcat output.
The _ScopedAllocatorConcat kernel outputs the backing tensor after performing runtime bounds checks. However, the shape of the backing tensor may not match the desired output shape of the concat operation. This change adds a "reshape" boolean attribute to _ScopedAllocatorConcat kernel. When this attribute is set to true, the kernel outputs a reshaped backing tensor according to the "shape" attribute. PiperOrigin-RevId: 195169105
Loading
Please sign in to comment