Home
last modified time | relevance | path

Searched refs:AssignDevice (Results 1 – 16 of 16) 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.cc166 p1->set_sharding(HloSharding::AssignDevice(1)); in TEST_F()
175 tuple_shape, {HloSharding::Tile(assignment), HloSharding::AssignDevice(1), in TEST_F()
181 ::testing::Not(op::Sharding(HloSharding::AssignDevice(1)))); in TEST_F()
184 ::testing::Not(op::Sharding(HloSharding::AssignDevice(0)))); in TEST_F()
185 EXPECT_THAT(p1.get(), op::Sharding(HloSharding::AssignDevice(1))); in TEST_F()
191 EXPECT_THAT(Explain(p0.get(), op::Sharding(HloSharding::AssignDevice(1))), in TEST_F()
197 EXPECT_THAT(Explain(p1.get(), op::Sharding(HloSharding::AssignDevice(0))), in TEST_F()
Dbatchnorm_expander.cc278 ? HloSharding::AssignDevice(unique_device.value()) in HandleBatchNormTraining()
365 ? HloSharding::AssignDevice(unique_device.value()) in HandleBatchNormInference()
542 ? 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.h1331 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.cc1373 OpSharding sharding = sharding_builder::AssignDevice(0); in Infeed()
1391 sharding_builder::AssignDevice(0); in Infeed()
/external/tensorflow/tensorflow/compiler/tf2xla/
Dsharding_util.cc76 xla::sharding_builder::AssignDevice(core)); in ParseShardingFromDevice()
Dxla_compiler.cc1025 it == arg_shardings.end() ? xla::sharding_builder::AssignDevice(0) in BuildArguments()
/external/tensorflow/tensorflow/core/common_runtime/
Dcolocation_graph.h84 Status AssignDevice(const Node& node);
Dcolocation_graph.cc440 Status Member::AssignDevice(const Node& node) { in AssignDevice() function in tensorflow::Member
974 return root_member.AssignDevice(node); in LimitToAssignedDevice()
/external/tensorflow/tensorflow/compiler/xla/tests/
Dlocal_client_execute_test.cc794 builder.SetSharding(sharding_builder::AssignDevice(1)); in XLA_TEST_F()