Home
last modified time | relevance | path

Searched refs:AssignDevice (Results 1 – 15 of 15) sorted by relevance

/external/tensorflow/tensorflow/compiler/xla/service/
Dhlo_sharding_test.cc57 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 …]
Dhlo_domain_test.cc432 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 …]
Dhlo_matchers_test.cc151 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()
Dbatchnorm_expander.cc311 ? HloSharding::AssignDevice(unique_device.value()) in HandleBatchNormTraining()
415 ? HloSharding::AssignDevice(unique_device.value()) in HandleBatchNormInference()
593 ? HloSharding::AssignDevice(unique_device.value()) in HandleBatchNormGrad()
Dhlo_sharding.h49 static HloSharding AssignDevice(int64 device_id);
Dhlo_sharding_metadata.cc253 instruction->shape(), HloSharding::AssignDevice(kUnassignedDevice)); in ApplyShardingFromUsers()
Dhlo_sharding.cc29 HloSharding HloSharding::AssignDevice(int64 device_id) { in AssignDevice() function in xla::HloSharding
Dhlo_instruction.h1167 set_single_sharding(HloSharding::AssignDevice(device)); in set_device_sharding()
/external/tensorflow/tensorflow/compiler/xla/client/
Dsharding_builder.h37 OpSharding AssignDevice(int device);
Dsharding_builder.cc27 OpSharding AssignDevice(int device) { in AssignDevice() function
Dxla_builder.cc1315 OpSharding sharding = sharding_builder::AssignDevice(0); in Infeed()
1334 sharding_builder::AssignDevice(0); in Infeed()
/external/tensorflow/tensorflow/compiler/tf2xla/
Dsharding_util.cc76 xla::sharding_builder::AssignDevice(core)); in ParseShardingFromDevice()
Dxla_compiler.cc213 : 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/
Dcolocation_graph.h74 Status AssignDevice(const Node& node, bool allow_soft_placement);
Dcolocation_graph.cc406 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()