Home
last modified time | relevance | path

Searched defs:layout_specs (Results 1 – 5 of 5) sorted by relevance

/aosp_15_r20/external/tensorflow/tensorflow/dtensor/mlir/expansions/
H A Din_top_k_spmd_expander.cc38 std::vector<ShardingSpec> layout_specs(2); in GetSuggestedPredictionsLayout() local
49 std::vector<ShardingSpec> layout_specs(layout.rank()); in MatchBatchDim() local
145 std::vector<std::string> layout_specs(1); in ComputeLayoutForward() local
H A Dsqueeze_spmd_expander.cc58 std::vector<ShardingSpec> layout_specs; in ComputeLayoutForward() local
91 std::vector<ShardingSpec> layout_specs; in ComputeLayoutBackward() local
H A Dsoftmax_spmd_expander.cc641 std::vector<ShardingSpec> layout_specs(2); in ComputeLayoutForward() local
690 std::vector<ShardingSpec> layout_specs(2); in ComputeLayoutBackward() local
H A Dtop_k_spmd_expander.cc32 std::vector<ShardingSpec> layout_specs(input_layout.rank()); in GetSuggestedLayout() local
H A Dmeta_spmd_expander.cc639 std::vector<std::string> layout_specs; in MakeLayoutForReshape() local