Home
last modified time | relevance | path

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

/aosp_15_r20/external/tensorflow/tensorflow/dtensor/mlir/
H A Dcollectives.cc366 const int mesh_dim_index = mesh.GetMeshDimIndexWithName(mesh_dim_name); in CreateConstSrcTargetPair() local
/aosp_15_r20/external/tensorflow/tensorflow/dtensor/mlir/expansions/
H A Dsoftmax_spmd_expander.cc269 const int mesh_dim_index = desired_layout.mesh().GetMeshDimIndexWithName( in ComputeOneHot() local
H A Dmeta_spmd_expander.cc1041 const int mesh_dim_index = in ExpandOp() local
/aosp_15_r20/external/tensorflow/tensorflow/dtensor/mlir/utils/
H A Dcollective_lowering.cc677 int mesh_dim_index = desired_layout.mesh().GetMeshDimIndexWithName( in LowerAllScatterOp() local