Home
last modified time | relevance | path

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

/external/tensorflow/tensorflow/dtensor/mlir/expansions/
Dmeta_spmd_expander.cc570 llvm::SmallVectorImpl<int64_t>& output_segment_start, in ComputeReshapeSegments() argument
592 output_segment_start.emplace_back(output_offset); in ComputeReshapeSegments()
638 llvm::SmallVectorImpl<int64_t>& output_segment_start) { in MakeLayoutForReshape() argument
650 if (output_shape[output_segment_start[i]] % num_shards == 0) in MakeLayoutForReshape()
651 layout_specs[output_segment_start[i]] = in MakeLayoutForReshape()
681 llvm::SmallVector<int64_t, 4> output_segment_start; in ExpandOp() local
691 output_segment_start, output_segment_end); in ExpandOp()
716 const int output_start = output_segment_start[i]; in ExpandOp()
817 llvm::SmallVector<int64_t, 4> output_segment_start; in ComputeLayoutForward() local
824 input_segment_end, output_segment_start, in ComputeLayoutForward()
[all …]