Home
last modified time | relevance | path

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

/external/tensorflow/tensorflow/compiler/xla/service/spmd/
Dspmd_partitioner.cc274 std::vector<int64> after_dim_size_stack; in AddInstruction() local
279 after_dim_size_stack.push_back(hlo->shape().dimensions(i)); in AddInstruction()
281 while (!before_dim_size_stack.empty() && !after_dim_size_stack.empty()) { in AddInstruction()
283 int64 after_size = after_dim_size_stack.back(); in AddInstruction()
287 hlo->shape().rank() - after_dim_size_stack.size(); in AddInstruction()
289 after_dim_size_stack.pop_back(); in AddInstruction()
301 after_dim_size_stack.push_back(after_size / before_size); in AddInstruction()
310 if (!before_dim_size_stack.empty() || !after_dim_size_stack.empty()) { in AddInstruction()