| /external/tensorflow/tensorflow/compiler/xla/service/ |
| D | hlo_sharding_test.cc | 72 HloSharding sharding = HloSharding::AssignDevice(5); in TEST_F()
156 HloSharding::AssignDevice(0)); in TEST_F()
172 *proto.add_tuple_shardings() = HloSharding::AssignDevice(0).ToProto(); in TEST_F()
180 EXPECT_EQ(shape_tree.element({1, 0}), HloSharding::AssignDevice(0)); in TEST_F()
207 HloSharding sharding1 = HloSharding::AssignDevice(1); in TEST_F()
208 HloSharding sharding2 = HloSharding::AssignDevice(1); in TEST_F()
213 HloSharding sharding1 = HloSharding::AssignDevice(1); in TEST_F()
214 HloSharding sharding2 = HloSharding::AssignDevice(2); in TEST_F()
249 *shape_tree2.mutable_element({0}) = HloSharding::AssignDevice(0); in TEST_F()
259 *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 | 166 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()
|
| D | batchnorm_expander.cc | 278 ? HloSharding::AssignDevice(unique_device.value()) in HandleBatchNormTraining()
365 ? HloSharding::AssignDevice(unique_device.value()) in HandleBatchNormInference()
542 ? HloSharding::AssignDevice(unique_device.value()) in HandleBatchNormGrad()
|
| D | hlo_sharding.h | 56 static HloSharding AssignDevice(int64 device_id,
|
| D | hlo_sharding_util_test.cc | 39 HloSharding sharding = HloSharding::AssignDevice(7); in TEST()
|
| D | hlo_sharding_metadata.cc | 253 instruction->shape(), HloSharding::AssignDevice(kUnassignedDevice)); in ApplyShardingFromUsers()
|
| D | hlo_sharding_util.cc | 614 return HloSharding::AssignDevice(*hlo.sharding().tile_assignment().begin(), in GatherEffectiveOutputSharding()
732 return HloSharding::AssignDevice(*index_sharding.tile_assignment().begin(), in ScatterEffectiveIndexSharding()
777 return HloSharding::AssignDevice(*data_sharding.tile_assignment().begin(), in ScatterEffectiveDataSharding()
|
| D | hlo_instruction_test.cc | 753 HloSharding sharding = HloSharding::AssignDevice(5); in TEST_F()
773 HloSharding sharding = HloSharding::AssignDevice(5); in TEST_F()
793 HloSharding sharding = HloSharding::AssignDevice(5); in TEST_F()
|
| D | hlo_sharding.cc | 31 HloSharding HloSharding::AssignDevice(int64 device_id, in AssignDevice() function in xla::HloSharding
|
| D | hlo_instruction.h | 1410 set_single_sharding(HloSharding::AssignDevice(device)); in set_device_sharding()
|
| D | sharding_propagation.cc | 1197 return HloSharding::AssignDevice(*new_tile_assignment.begin(), in GetShardingFromUser()
|
| /external/tensorflow/tensorflow/compiler/xla/client/ |
| D | sharding_builder.h | 40 OpSharding AssignDevice(int device);
|
| D | sharding_builder.cc | 33 OpSharding AssignDevice(int device) { in AssignDevice() function
|
| D | xla_builder.cc | 1696 OpSharding sharding = sharding_builder::AssignDevice(0); in Infeed()
1714 sharding_builder::AssignDevice(0); in Infeed()
|
| /external/tensorflow/tensorflow/core/common_runtime/ |
| D | colocation_graph.h | 88 Status AssignDevice(const Node& node);
|
| D | colocation_graph.cc | 467 Status Member::AssignDevice(const Node& node) { in AssignDevice() function in tensorflow::Member
1103 return root_member.AssignDevice(node); in LimitToAssignedDevice()
|
| /external/tensorflow/tensorflow/compiler/tf2xla/ |
| D | sharding_util.cc | 80 auto sharding = xla::sharding_builder::AssignDevice(core); in ParseShardingFromDevice()
|
| D | xla_compiler.cc | 1061 it == arg_shardings.end() ? xla::sharding_builder::AssignDevice(0) in BuildArguments()
|
| /external/tensorflow/tensorflow/compiler/mlir/tensorflow/transforms/ |
| D | tpu_sharding_identification_pass.cc | 313 xla::sharding_builder::AssignDevice(0).SerializeAsString(); in IdentifyXlaShardingForTPUComputation()
|
| /external/tensorflow/tensorflow/compiler/xla/service/spmd/ |
| D | spmd_partitioner.cc | 1420 subsharding = HloSharding::AssignDevice(0); in Preprocess()
1426 return HloSharding::AssignDevice(0); in Preprocess()
2181 const HloSharding sharding = HloSharding::AssignDevice(device); in HandleSingleDevice()
3054 .Reshard(HloSharding::AssignDevice(0)) in HandleRng()
3059 clone->set_sharding(HloSharding::AssignDevice(0)); in HandleRng()
3092 rng->set_sharding(HloSharding::AssignDevice(0)); in HandleRng()
3757 hlo->set_sharding(HloSharding::AssignDevice(0)); in PreprocessSharding()
|
| /external/tensorflow/tensorflow/core/tpu/kernels/xla/ |
| D | host_compute_ops.cc | 136 const auto device_sharding = xla::sharding_builder::AssignDevice(tpu_core_); in Compile()
|
| /external/tensorflow/tensorflow/compiler/mlir/xla/transforms/ |
| D | legalize_tf_communication.cc | 216 ::xla::sharding_builder::AssignDevice(tpu_core).SerializeAsString(); in SetOpSharding()
|
| /external/tensorflow/tensorflow/compiler/xla/tests/ |
| D | local_client_execute_test.cc | 794 builder.SetSharding(sharding_builder::AssignDevice(1)); in XLA_TEST_F()
|
| /external/tensorflow/tensorflow/core/tpu/graph_rewrite/ |
| D | distributed_tpu_rewrite_pass.cc | 1232 auto core_zero_sharding = xla::sharding_builder::AssignDevice(0); in ParseAndValidateSharding()
2064 xla::sharding_builder::AssignDevice(*assigned_core)); in AssignArgsAndRetvalsToCores()
2177 xla::sharding_builder::AssignDevice(*assigned_core)); in AssignArgsAndRetvalsToCores()
|