Searched defs:branch_index (Results 1 – 11 of 11) sorted by relevance
54 int branch_index = 0; in TryRemoveConditional() local
2556 const Shape& branch_index, in InferConditionalShape()
1264 int branch_index; in HandleConditional() local
970 const Shape& shape, HloInstruction* branch_index, in CreateConditional()
64 int32 branch_index = -1; in ExecuteOnStream() local
428 int branch_index = static_cast<int>(branch); in BuildArgumentNodes() local440 int branch_index = e->src_output(); in BuildArgumentNodes() local551 int branch_index = static_cast<int>(branch); in ExtractBodies() local693 int branch_index = static_cast<int>(branch); in ExtractBodies() local705 int branch_index = static_cast<int>(find_branch(e)); in ExtractBodies() local728 int branch_index = static_cast<int>(branch); in BuildIfNode() local893 int branch_index = static_cast<int>(branch); in BuildAndReplace() local
197 XlaOp branch_index; in XLA_TEST_P() local240 XlaOp branch_index; in XLA_TEST_P() local415 XlaOp branch_index; in XLA_TEST_P() local
471 Status MakeCurrentIterator(IteratorContext* ctx, int64 branch_index, in MakeCurrentIterator()
243 const Tensor& branch_index = ctx->input(0); in ComputeAsync() local
1902 const XlaOp& branch_index, in Conditional()3411 XlaOp Conditional(const XlaOp& branch_index, in Conditional()
2532 auto branch_index = conditional->operand(0); in HandleConditional() local