Make root determination of XLA computations in XlaBuilder less magical.
Previously (before this CL), in the XLA builder certain XLA ops could never be the root of a computation. These restricted ops include Send, Outfeed, and several others. The root of the built computation was then the last added op which was not in this restricted set. However, this is undesirable because now Send and Outfeed produce token values and it may be desirable to return these tokens from a computation, something which is impossible now. This CL addresses this problem by allowing any op to be the root of a computation. This means now the root of the computation will be the last operation added before calling Build(). Furthermore, to enable previous functionality and improve expressiveness in general a new XlaComputation::Build method is added which takes an XlaOp which specifies the root. PiperOrigin-RevId: 207887842
Loading
Please sign in to comment