Searched refs:AssignDevice (Results 1 – 15 of 15) sorted by relevance
/external/tensorflow/tensorflow/compiler/xla/service/ |
D | hlo_sharding_test.cc | 57 HloSharding sharding = HloSharding::AssignDevice(5); in TEST_F() 121 HloSharding::AssignDevice(0)); in TEST_F() 137 *proto.add_tuple_shardings() = HloSharding::AssignDevice(0).ToProto(); in TEST_F() 145 EXPECT_EQ(shape_tree.element({1, 0}), HloSharding::AssignDevice(0)); in TEST_F() 172 HloSharding sharding1 = HloSharding::AssignDevice(1); in TEST_F() 173 HloSharding sharding2 = HloSharding::AssignDevice(1); in TEST_F() 178 HloSharding sharding1 = HloSharding::AssignDevice(1); in TEST_F() 179 HloSharding sharding2 = HloSharding::AssignDevice(2); in TEST_F() 214 *shape_tree2.mutable_element({0}) = HloSharding::AssignDevice(0); in TEST_F() 224 *shape_tree1.mutable_element({0}) = HloSharding::AssignDevice(0); in TEST_F() [all …]
|
D | hlo_domain_test.cc | 432 HloSharding::Tuple(new_tuple->shape(), {HloSharding::AssignDevice(1), in TEST_F() 433 HloSharding::AssignDevice(0)})); in TEST_F() 472 EXPECT_EQ(root->sharding(), HloSharding::AssignDevice(1)); in TEST_F() 527 EXPECT_EQ(HloSharding::Tuple(tpl->shape(), {HloSharding::AssignDevice(1), in TEST_F() 528 HloSharding::AssignDevice(0)}), in TEST_F() 665 EXPECT_EQ(HloSharding::Tuple(tuple0->shape(), {HloSharding::AssignDevice(1), in TEST_F() 666 HloSharding::AssignDevice(1), in TEST_F() 667 HloSharding::AssignDevice(0)}), in TEST_F() 671 EXPECT_EQ(HloSharding::Tuple(copy0->shape(), {HloSharding::AssignDevice(1), in TEST_F() 672 HloSharding::AssignDevice(0)}), in TEST_F() [all …]
|
D | hlo_matchers_test.cc | 151 p1->set_sharding(HloSharding::AssignDevice(1)); in TEST() 160 tuple_shape, {HloSharding::Tile(assignment), HloSharding::AssignDevice(1), in TEST() 166 ::testing::Not(op::Sharding(HloSharding::AssignDevice(1)))); in TEST() 169 ::testing::Not(op::Sharding(HloSharding::AssignDevice(0)))); in TEST() 170 EXPECT_THAT(p1.get(), op::Sharding(HloSharding::AssignDevice(1))); in TEST() 176 EXPECT_THAT(Explain(p0.get(), op::Sharding(HloSharding::AssignDevice(1))), in TEST() 182 EXPECT_THAT(Explain(p1.get(), op::Sharding(HloSharding::AssignDevice(0))), in TEST()
|
D | batchnorm_expander.cc | 311 ? HloSharding::AssignDevice(unique_device.value()) in HandleBatchNormTraining() 415 ? HloSharding::AssignDevice(unique_device.value()) in HandleBatchNormInference() 593 ? HloSharding::AssignDevice(unique_device.value()) in HandleBatchNormGrad()
|
D | hlo_sharding.h | 49 static HloSharding AssignDevice(int64 device_id);
|
D | hlo_sharding_metadata.cc | 253 instruction->shape(), HloSharding::AssignDevice(kUnassignedDevice)); in ApplyShardingFromUsers()
|
D | hlo_sharding.cc | 29 HloSharding HloSharding::AssignDevice(int64 device_id) { in AssignDevice() function in xla::HloSharding
|
D | hlo_instruction.h | 1167 set_single_sharding(HloSharding::AssignDevice(device)); in set_device_sharding()
|
/external/tensorflow/tensorflow/compiler/xla/client/ |
D | sharding_builder.h | 37 OpSharding AssignDevice(int device);
|
D | sharding_builder.cc | 27 OpSharding AssignDevice(int device) { in AssignDevice() function
|
D | xla_builder.cc | 1315 OpSharding sharding = sharding_builder::AssignDevice(0); in Infeed() 1334 sharding_builder::AssignDevice(0); in Infeed()
|
/external/tensorflow/tensorflow/compiler/tf2xla/ |
D | sharding_util.cc | 76 xla::sharding_builder::AssignDevice(core)); in ParseShardingFromDevice()
|
D | xla_compiler.cc | 213 : xla::sharding_builder::AssignDevice(it->second)); in BuildComputation() 285 : xla::sharding_builder::AssignDevice(core)); in BuildComputation() 798 xla::sharding_builder::AssignDevice(core); in BuildArguments() 823 : xla::sharding_builder::AssignDevice(core)); in BuildArguments() 832 : xla::sharding_builder::AssignDevice(core)); in BuildArguments()
|
/external/tensorflow/tensorflow/core/common_runtime/ |
D | colocation_graph.h | 74 Status AssignDevice(const Node& node, bool allow_soft_placement);
|
D | colocation_graph.cc | 406 Status Member::AssignDevice(const Node& node, bool allow_soft_placement) { in AssignDevice() function in tensorflow::Member 682 return root_member.AssignDevice(node, allow_soft_placement_); in LimitToAssignedDevice()
|