joel
joel

Reputation: 7867

How to compose XlaComputation?

How do I use an XlaComputation when building another XlaComputation? For example, I want to build Add into a computation, which I'm trying to do as

xla::XlaBuilder add_builder("Add");
auto one = ConstantR0(&add_builder, 1);
auto two = ConstantR0(&add_builder, 2);
auto sum = Add(one, two);
add_builder.Build();

xla::XlaBuilder mul_builder("Mul");
auto three = ConstantR0(&mul_builder, 3);
auto res = Mul(sum, three);
mul_builder.Build();

but I get

Attempting to fetch value instead of handling error Invalid argument: XlaOp with handle 3 is built by builder 'Add', but is trying to use it in builder 'Mul':

I'm aware I'm not supposed to use an XlaOp from one builder in another computation, but I don't understand how to avoid that when that XlaOp corresponds to the result of that computation. After reading the following comment in the source code

// This represents an instruction that has been enqueued using the XlaBuilder.
// This is used to pass to subsequent computations that depends upon the
// instruction as an operand.
class XlaOp {

I'm guessing that only applies to XlaOps using the same builder.

Upvotes: 1

Views: 110

Answers (0)

Related Questions