Home
last modified time | relevance | path

Searched defs:desired_layout (Results 1 – 4 of 4) sorted by relevance

/aosp_15_r20/external/tensorflow/tensorflow/dtensor/mlir/
H A Ddtensor_allreduce_scatter_optimization.cc52 const Layout desired_layout = all_scatter.output_layout(); in GetScatterGroupAssignment() local
88 const Layout desired_layout = all_scatter.output_layout(); in ApplyOptimization() local
H A Dcollectives.cc107 const Layout& original_layout, const Layout& desired_layout, in EmitAllScatter()
/aosp_15_r20/external/tensorflow/tensorflow/dtensor/mlir/expansions/
H A Dsoftmax_spmd_expander.cc231 const Layout& desired_layout) { in ComputeOneHot()
/aosp_15_r20/external/tensorflow/tensorflow/dtensor/mlir/utils/
H A Dcollective_lowering.cc632 const Layout desired_layout = all_scatter.output_layout(); in LowerAllScatterOp() local