Searched refs:token_output (Results 1 – 4 of 4) sorted by relevance
247 xla::XlaOp token_output = in Compile() local249 auto shape_or = b->GetShape(token_output); in Compile()255 OP_REQUIRES_OK(ctx, compiler->SetNodeToken(name(), token_output)); in Compile()
259 xla::XlaOp token_output = in Compile() local261 auto shape_or = b->GetShape(token_output); in Compile()267 OP_REQUIRES_OK(ctx, compiler->SetNodeToken(name(), token_output)); in Compile()
541 xla::XlaOp token_output = in Compile() local543 auto shape_or = builder->GetShape(token_output); in Compile()549 OP_REQUIRES_OK(ctx, compiler->SetNodeToken(name(), token_output)); in Compile()
162 std::unique_ptr<xla::XlaOp> token_output, in BuildComputation() argument303 if (token_output) { in BuildComputation()304 elems.push_back(*token_output); in BuildComputation()1074 std::unique_ptr<xla::XlaOp> token_output; in CompileGraph() local1137 token_output.reset(new xla::XlaOp(xla::AfterAll(&builder, token_inputs))); in CompileGraph()1163 std::move(token_output), in CompileGraph()