Home
last modified time | relevance | path

Searched refs:body_root (Results 1 – 8 of 8) sorted by relevance

/external/tensorflow/tensorflow/compiler/xla/service/
Dbfloat16_propagation.cc175 auto body_root = body->root_instruction(); in DetermineWhileComputationsPrecision() local
179 body_root->shape(), [this, while_hlo, body_root]( in DetermineWhileComputationsPrecision()
185 AddToOrRemoveFromBF16ChangeSet(body_root, index, BF16); in DetermineWhileComputationsPrecision()
186 VLOG(2) << "While body root " << body_root->ToString() in DetermineWhileComputationsPrecision()
Ddynamic_dimension_inference.cc1509 HloInstruction* body_root = hlo->while_body()->root_instruction(); in HandleWhile() local
1510 std::vector<HloInstruction*> new_root_operands(body_root->operand_count(), in HandleWhile()
1517 body_root->shape().tuple_shapes(i), body_root, i)); in HandleWhile()
Dsharding_propagation.cc1619 auto body_root = while_body->root_instruction(); in CheckAndUpdateDeviceAssignmentsInWhileBody() local
1620 auto root_device = body_root->sharding_unique_device(); in CheckAndUpdateDeviceAssignmentsInWhileBody()
1622 body_root->set_device_sharding(*unique_device); in CheckAndUpdateDeviceAssignmentsInWhileBody()
1624 return bad_status(body_root, *root_device, channel_instruction, in CheckAndUpdateDeviceAssignmentsInWhileBody()
Dhlo_alias_analysis_test.cc1215 auto body_root = body_builder.AddInstruction( in TEST_F() local
1242 schedule.set_sequence(body, {body_param, body_root}); in TEST_F()
Dbuffer_assignment_test.cc1067 HloInstruction* body_root = body_computation->root_instruction(); in TEST_F() local
1068 EXPECT_EQ(HloOpcode::kTuple, body_root->opcode()); in TEST_F()
1074 [this, &buffers, while_op, body_root](const Shape& /*subshape*/, in TEST_F()
1077 auto body_root_allocation = GetAllocation(*buffers, body_root, index); in TEST_F()
Dsharding_propagation_test.cc1698 auto body_root = FindInstruction(module.get(), "tuple"); in TEST_P() local
1699 EXPECT_NE(nullptr, body_root); in TEST_P()
1704 body_root->set_sharding(sharding); in TEST_P()
Dhlo_dataflow_analysis_test.cc402 auto body_root = body_builder.AddInstruction( in TEST_P() local
449 UnorderedElementsAre(HloUse{add, 0, {}}, HloUse{body_root, 0, {}}, in TEST_P()
/external/tensorflow/tensorflow/compiler/xla/service/spmd/
Ddot_handler.cc3180 auto body_root = body->root_instruction(); in MoveUsersIntoWindowedDotGeneralLoopOnNonContractingDimensions() local
3181 CHECK_EQ(body_root->opcode(), HloOpcode::kTuple); in MoveUsersIntoWindowedDotGeneralLoopOnNonContractingDimensions()
3185 body_param, body_root}) { in MoveUsersIntoWindowedDotGeneralLoopOnNonContractingDimensions()
3214 auto dus = body_root->mutable_operand(2); in MoveUsersIntoWindowedDotGeneralLoopOnNonContractingDimensions()
3371 body_root->ReplaceOperandWithDifferentShape(2, new_output_inside)); in MoveUsersIntoWindowedDotGeneralLoopOnNonContractingDimensions()