Searched defs:splitAxes (Results 1 – 2 of 2) sorted by relevance
397 SmallVector<SmallVector<MeshAxis>> splitAxes(resultType.getRank()); getShardingAttribute() local 439 SmallVector<SmallVector<MeshAxis>> splitAxes(operandType.getRank()); getShardingAttribute() local
153 shardShape(const InShape & inShape,const MeshShape & meshShape,const SplitAxes & splitAxes,OutShape & outShape) shardShape() argument 339 verify(function_ref<InFlightDiagnostic ()> emitError,FlatSymbolRefAttr,ArrayRef<MeshAxesAttr> splitAxes,ArrayRef<MeshAxis> partialAxes,ReductionKind) verify() argument