Searched refs:gather_slice_sizes (Results 1 – 16 of 16) sorted by relevance
| /external/tensorflow/tensorflow/compiler/xla/service/ |
| D | gather_expander.cc | 205 gather.gather_slice_sizes())); in GatherLoopBody()
291 return absl::c_equal(gather_instr->gather_slice_sizes(), in GatherIsBroadcast()
379 gather_instr->gather_slice_sizes(), gather_loop_trip_count, in ExpandInstruction()
415 absl::c_equal(inst->gather_slice_sizes(), in InstructionMatchesPattern()
|
| D | gather_simplifier.cc | 38 if (absl::c_linear_search(gather->gather_slice_sizes(), 0)) { in ExpandInstruction()
61 auto slice_sizes = Permute(gather->gather_slice_sizes(), operand_permutation); in ExpandInstruction()
|
| D | hlo_instruction.cc | 843 std::vector<int64_t> gather_slice_sizes; in CreateFromProto() local
844 const auto& slice_sizes = proto.gather_slice_sizes(); in CreateFromProto()
845 gather_slice_sizes.reserve(slice_sizes.size()); in CreateFromProto()
847 gather_slice_sizes.push_back(bound); in CreateFromProto()
850 *gather_dimension_numbers, gather_slice_sizes, in CreateFromProto()
4822 absl::Span<const int64_t> HloInstruction::gather_slice_sizes() const { in gather_slice_sizes() function in xla::HloInstruction
4823 return Cast<HloGatherInstruction>(this)->gather_slice_sizes(); in gather_slice_sizes()
|
| D | dynamic_padder.cc | 738 std::vector<int64_t> gather_slice_sizes(output_shape.dimensions().begin(), in RewriteDynamicReshapeCombineInput() local
740 gather_slice_sizes[output_dim] = 1; in RewriteDynamicReshapeCombineInput()
743 gather_slice_sizes, true)); in RewriteDynamicReshapeCombineInput()
|
| D | hlo_instructions.cc | 3101 for (int64_t bound : gather_slice_sizes()) { in ToProto()
3112 StrCat("slice_sizes={", StrJoin(gather_slice_sizes(), ","), "}")}; in ExtraAttributesToStringImpl()
3127 gather_slice_sizes() == casted_other.gather_slice_sizes() && in IdenticalSlowPath()
3137 gather_slice_sizes(), indices_are_sorted()); in CloneWithNewOperandsImpl()
|
| D | dynamic_dimension_inference.cc | 1372 if (hlo->gather_slice_sizes()[input_dynamic_dimension] == 1) { in HandleGather()
1377 if (hlo->gather_slice_sizes()[input_dynamic_dimension] == in HandleGather()
|
| D | indexed_array_analysis.cc | 161 instr->gather_slice_sizes(), in ComputeArrayFor()
|
| D | hlo_sharding_util.cc | 1098 start_index_map, offset_dims, hlo.gather_slice_sizes()); in GatherDataOperandShardingFromOutput()
1545 if (hlo.gather_slice_sizes()[dnums.start_index_map(i)] == 1) { in GetGatherBatchParallelDims()
|
| D | hlo.proto | 213 repeated int64 gather_slice_sizes = 34; field
|
| D | hlo_instruction.h | 2105 absl::Span<const int64_t> gather_slice_sizes() const;
|
| D | hlo_instructions.h | 2110 absl::Span<const int64_t> gather_slice_sizes() const { in gather_slice_sizes() function
|
| D | sharding_propagation.cc | 2164 instruction->gather_slice_sizes(), instruction->shape(), in InferShardingFromOperands()
|
| D | hlo_verifier.cc | 1659 gather->gather_dimension_numbers(), gather->gather_slice_sizes())); in HandleGather()
|
| D | hlo_evaluator.cc | 2484 shape.dimensions_size(), gather->gather_slice_sizes(), dim_numbers); in HandleGather()
|
| /external/tensorflow/tensorflow/compiler/xla/service/spmd/ |
| D | gather_scatter_handler.cc | 385 operand, start_index_map, gather->gather_slice_sizes()); in PartitionGatherTrivialIndexedOperandDimension()
700 gather->gather_slice_sizes(), this)); in HandleGather()
|
| /external/tensorflow/tensorflow/compiler/mlir/xla/ |
| D | hlo_function_importer.cc | 699 gather_instruction->gather_slice_sizes().begin(), in ImportInstructionImpl()
700 gather_instruction->gather_slice_sizes().end()); in ImportInstructionImpl()
|