Home
last modified time | relevance | path

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

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