Home
last modified time | relevance | path

Searched defs:all_gather_dimension (Results 1 – 7 of 7) sorted by relevance

/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/
H A Dall_gather_combiner.cc56 int64_t all_gather_dimension = in CombineAllGathers() local
H A Dshape_inference.cc2036 absl::Span<const Shape* const> operand_shapes, int64_t all_gather_dimension, in InferAllGatherShape()
2060 absl::Span<const Shape* const> operand_shapes, int64_t all_gather_dimension, in InferAllGatherStartShape()
H A Dhlo_instruction.cc478 int64_t all_gather_dimension = proto.dimensions(0); in CreateFromProto() local
1259 int64_t all_gather_dimension, absl::Span<const ReplicaGroup> replica_groups, in CreateAllGather()
1270 int64_t all_gather_dimension, absl::Span<const ReplicaGroup> replica_groups, in CreateAllGatherStart()
H A Dhlo_instructions.h586 int64_t all_gather_dimension() const { return all_gather_dimension_; } in all_gather_dimension() function
H A Dhlo_instructions.cc787 absl::Span<HloInstruction* const> operands, int64_t all_gather_dimension, in HloAllGatherInstruction()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/spmd/
H A Dspmd_partitioner_util.cc1545 int64_t channel_id, int64_t all_gather_dimension) { in GetPerGroupCollectiveOpsCreator()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/client/
H A Dxla_builder.cc2868 XlaOp XlaBuilder::AllGather(XlaOp operand, int64_t all_gather_dimension, in AllGather()
4653 XlaOp AllGather(const XlaOp operand, int64_t all_gather_dimension, in AllGather()