Home
last modified time | relevance | path

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

/aosp_15_r20/external/tensorflow/tensorflow/dtensor/mlir/expansions/
H A Ddataparallel_spmd_expander.cc161 std::vector<std::string> batch_specs(batch_rank, Layout::kUnshardedDim); in IntermediateBatchLayout() local
189 batch_specs[i] = mesh_dim; in IntermediateBatchLayout()
193 for (const std::string& spec : batch_specs) counter[spec] += 1; in IntermediateBatchLayout()
194 for (std::string& spec : batch_specs) { in IntermediateBatchLayout()
200 return Layout::GetLayout(batch_specs, mesh); in IntermediateBatchLayout()