Commit 7b0760fe authored by Mark Heffernan's avatar Mark Heffernan Committed by TensorFlower Gardener
Browse files

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
parent a26e1ef0
Loading
Loading
Loading
Loading
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please to comment