Lines Matching defs:MeshAxis
44 using MeshAxis = mesh::MeshAxis;
133 ArrayRef<MeshAxis> reductionMeshAxes, MeshOp meshOp,
179 ArrayRef<MeshAxis> reductionMeshAxes, IRMapping &spmdizationMap,
194 Value unshardedLinalgOpResult, ArrayRef<MeshAxis> opReductionMeshAxes,
197 SmallVector<MeshAxis> allReduceMeshAxes;
199 [&resultSharding](MeshAxis axis) {
215 LinalgOp unshardedOp, ArrayRef<MeshAxis> opReductionMeshAxes,
232 ArrayRef<SmallVector<MeshAxis>> meshAxisAssignmentForLoopIterators,
236 SmallVector<MeshAxis> reductionMeshAxes = mesh::getReductionMeshAxes(