Home
last modified time | relevance | path

Searched refs:ReplicaGroup (Results 1 – 25 of 33) sorted by relevance

12

/external/tensorflow/tensorflow/core/tpu/kernels/
Dcross_replica_ops.cc28 std::vector<xla::ReplicaGroup> Convert( in Convert()
34 std::vector<xla::ReplicaGroup> replica_groups; in Convert()
40 xla::ReplicaGroup replica_group; in Convert()
59 std::vector<xla::ReplicaGroup> replica_groups = in Compile()
81 std::vector<xla::ReplicaGroup> replica_groups = in Compile()
/external/tensorflow/tensorflow/compiler/xla/service/
Dcollective_ops_utils.cc55 absl::Span<const ReplicaGroup> replica_groups) { in GetParticipatingReplicas()
64 absl::optional<ReplicaGroup> replica_group; in GetParticipatingReplicas()
65 for (const ReplicaGroup& g : replica_groups) { in GetParticipatingReplicas()
80 int total_replica_count, absl::Span<const ReplicaGroup> replica_groups) { in GetParticipatingDevices()
Dcollective_ops_utils_test.cc36 std::vector<ReplicaGroup> replica_groups(3); in TEST()
76 std::vector<ReplicaGroup> replica_groups(2); in TEST()
Dhlo_parser.h83 StatusOr<std::vector<ReplicaGroup>> ParseReplicaGroupsOnly(
Dcollective_ops_utils.h44 absl::Span<const ReplicaGroup> replica_groups);
50 int total_replica_count, absl::Span<const ReplicaGroup> replica_groups);
Dall_reduce_combiner_test.cc101 std::vector<ReplicaGroup> CreateReplicaGroups( in CreateReplicaGroups()
103 std::vector<ReplicaGroup> replica_groups(groups.size()); in CreateReplicaGroups()
Dhlo_instruction.h635 const std::vector<ReplicaGroup>& replica_groups, bool constrain_layout,
654 const std::vector<ReplicaGroup>& replica_groups, bool constrain_layout,
687 const std::vector<ReplicaGroup>& replica_groups, bool constrain_layout,
1813 const std::vector<ReplicaGroup>& replica_groups() const;
2170 string ReplicaGroupsToString(const std::vector<ReplicaGroup>& replica_groups);
Dhlo_instructions.h349 const std::vector<ReplicaGroup>& replica_groups() const { in replica_groups()
372 const std::vector<ReplicaGroup>& replica_groups, bool constrain_layout,
384 std::vector<ReplicaGroup> replica_groups_;
392 const std::vector<ReplicaGroup>& replica_groups, bool constrain_layout,
425 const std::vector<ReplicaGroup>& replica_groups, bool constrain_layout,
467 const std::vector<ReplicaGroup>& replica_groups, bool constrain_layout,
Dbfloat16_normalization_test.cc287 std::vector<ReplicaGroup> replica_groups(1); in TEST_F()
316 std::vector<ReplicaGroup> replica_groups(1); in TEST_F()
Dhlo_parser.cc221 StatusOr<std::vector<ReplicaGroup>> ParseReplicaGroupsOnly();
417 bool ParseReplicaGroupsOnly(std::vector<ReplicaGroup>* replica_groups);
563 std::vector<ReplicaGroup> CreateReplicaGroups( in CreateReplicaGroups()
565 std::vector<ReplicaGroup> replica_groups; in CreateReplicaGroups()
568 ReplicaGroup group; in CreateReplicaGroups()
1215 std::vector<ReplicaGroup> replica_groups; in ParseInstructionRhs()
1244 std::vector<ReplicaGroup> replica_groups; in ParseInstructionRhs()
1270 std::vector<ReplicaGroup> replica_groups; in ParseInstructionRhs()
2887 std::vector<ReplicaGroup>* replica_groups) { in ParseReplicaGroupsOnly()
5136 StatusOr<std::vector<ReplicaGroup>> HloParserImpl::ParseReplicaGroupsOnly() { in ParseReplicaGroupsOnly()
[all …]
Dall_reduce_combiner.cc156 std::vector<ReplicaGroup> replica_groups;
Dhlo_instruction.cc425 std::vector<ReplicaGroup>(proto.replica_groups().begin(), in CreateFromProto()
446 std::vector<ReplicaGroup>(proto.replica_groups().begin(), in CreateFromProto()
470 std::vector<ReplicaGroup>(proto.replica_groups().begin(), in CreateFromProto()
1045 const std::vector<ReplicaGroup>& replica_groups, bool constrain_layout, in CreateAllGather()
1055 const std::vector<ReplicaGroup>& replica_groups, bool constrain_layout, in CreateAllReduce()
1064 const std::vector<ReplicaGroup>& replica_groups, bool constrain_layout, in CreateAllToAll()
3733 string ReplicaGroupsToString(const std::vector<ReplicaGroup>& replica_groups) { in ReplicaGroupsToString()
3736 for (const ReplicaGroup& group : replica_groups) { in ReplicaGroupsToString()
4127 const std::vector<ReplicaGroup>& HloInstruction::replica_groups() const { in replica_groups()
Dar_crs_combiner_test.cc378 void CompareReplicaGroups(const std::vector<ReplicaGroup>& groups_before, in CompareReplicaGroups()
379 const std::vector<ReplicaGroup>& groups_after) { in CompareReplicaGroups()
Dhlo_instructions.cc573 const std::vector<ReplicaGroup>& replica_groups, bool constrain_layout, in HloCollectiveInstruction()
613 [](const ReplicaGroup& a, const ReplicaGroup& b) { in IdenticalSlowPathIgnoringChannelIdValues()
620 const std::vector<ReplicaGroup>& replica_groups, bool constrain_layout, in HloAllGatherInstruction()
668 const std::vector<ReplicaGroup>& replica_groups, bool constrain_layout, in HloAllReduceInstruction()
724 const std::vector<ReplicaGroup>& replica_groups, bool constrain_layout, in HloAllToAllInstruction()
/external/tensorflow/tensorflow/compiler/mlir/xla/
Dattribute_exporter.cc90 StatusOr<std::vector<ReplicaGroup>> ConvertReplicaGroups( in ConvertReplicaGroups()
100 std::vector<ReplicaGroup> replica_groups(type.getDimSize(0)); in ConvertReplicaGroups()
101 for (ReplicaGroup& group : replica_groups) { in ConvertReplicaGroups()
Dattribute_exporter.h35 StatusOr<std::vector<ReplicaGroup>> ConvertReplicaGroups(
Dhlo_function_importer.h77 const std::vector<ReplicaGroup>& replica_groups, mlir::Builder* builder);
/external/tensorflow/tensorflow/core/profiler/protobuf/
Dpod_viewer.proto12 message ReplicaGroup { message
29 repeated ReplicaGroup replica_groups = 5;
/external/tensorflow/tensorflow/compiler/xla/service/gpu/
Dnccl_collective_thunk.h71 std::vector<ReplicaGroup> replica_groups;
/external/tensorflow/tensorflow/compiler/xla/client/
Dxla_builder.h731 absl::Span<const ReplicaGroup> replica_groups = {});
735 absl::Span<const ReplicaGroup> replica_groups = {},
742 absl::Span<const ReplicaGroup> replica_groups = {},
748 const std::vector<ReplicaGroup>& replica_groups,
753 const std::vector<ReplicaGroup>& replica_groups,
1296 absl::Span<const ReplicaGroup> replica_groups);
1299 absl::Span<const ReplicaGroup> replica_groups,
1304 absl::Span<const ReplicaGroup> replica_groups,
1309 const std::vector<ReplicaGroup>& replica_groups,
1313 const std::vector<ReplicaGroup>& replica_groups,
[all …]
Dxla_builder.cc2716 absl::Span<const ReplicaGroup> replica_groups, in AllGather()
2734 for (const ReplicaGroup& group : replica_groups) { in AllGather()
2752 XlaOp operand, absl::Span<const ReplicaGroup> replica_groups) { in CrossReplicaSum()
2782 absl::Span<const ReplicaGroup> replica_groups, in AllReduce()
2834 for (const ReplicaGroup& group : replica_groups) { in AllReduce()
2859 const std::vector<ReplicaGroup>& replica_groups, in AllToAll()
2873 int64 split_count, const std::vector<ReplicaGroup>& replica_groups) { in AllToAllArray()
2888 for (const ReplicaGroup& group : replica_groups) { in AllToAllArray()
2925 const std::vector<ReplicaGroup>& replica_groups, in AllToAllTuple()
2979 for (const ReplicaGroup& group : replica_groups) { in AllToAllTuple()
[all …]
/external/tensorflow/tensorflow/compiler/xla/python/
Dxla_client.py663 class ReplicaGroup(object): class
672 replica_group_proto = ReplicaGroup()
Dtypes.h348 struct type_caster<xla::ReplicaGroup> {
350 PYBIND11_TYPE_CASTER(xla::ReplicaGroup, _("xla::ReplicaGroup"));
Dops.cc68 XlaOp, const XlaComputation&, absl::Span<const ReplicaGroup>, in BuildOpsSubmodule()
83 static_cast<XlaOp (*)(XlaOp, absl::Span<const ReplicaGroup>)>( in BuildOpsSubmodule()
/external/tensorflow/tensorflow/compiler/xla/service/cpu/
Dcpu_runtime.cc597 std::vector<xla::ReplicaGroup> group, xla::int32 channel_id_present, in GetRendezvousKey()
630 std::vector<xla::ReplicaGroup> group = in __xla_cpu_runtime_AllToAll()
668 std::vector<xla::ReplicaGroup> group = in __xla_cpu_runtime_AllReduce()

12