Home
last modified time | relevance | path

Searched refs:input_segment_start (Results 1 – 1 of 1) sorted by relevance

/external/tensorflow/tensorflow/dtensor/mlir/expansions/
Dmeta_spmd_expander.cc568 llvm::SmallVectorImpl<int64_t>& input_segment_start, in ComputeReshapeSegments() argument
591 input_segment_start.emplace_back(input_offset); in ComputeReshapeSegments()
637 llvm::SmallVectorImpl<int64_t>& input_segment_start, in MakeLayoutForReshape() argument
647 for (int i = 0; i < input_segment_start.size(); ++i) { in MakeLayoutForReshape()
649 input_layout.dim(input_segment_start[i])); in MakeLayoutForReshape()
652 input_layout.sharding_spec(input_segment_start[i]); in MakeLayoutForReshape()
679 llvm::SmallVector<int64_t, 4> input_segment_start; in ExpandOp() local
690 input_segment_start, input_segment_end, in ExpandOp()
714 for (int i = 0; i < input_segment_start.size(); ++i) { in ExpandOp()
715 const int input_start = input_segment_start[i]; in ExpandOp()
[all …]