Lines Matching defs:axis

1071 // that reduces across the specified axis.
1072 static LogicalResult reduceMatchAndRewriteHelper(Operation *op, uint64_t axis,
1083     if (axis != i) {
1110       loc, input, filledTensor, axis,
1130     int32_t dimToPush = i > axis ? i + 1 : i;
1135     int32_t expandedDim = axis < expandInputRank ? axis : expandInputRank - 1;
1846     auto axis = op.getAxis();
1855     Value axisDimSize = rewriter.create<tensor::DimOp>(loc, input, axis);
1874             if (i == axis) {
1987     int axis = argmaxOp.getAxis();
1997       if (inputTy.isDynamicDim(i) && i != axis) {
2035     // We need to reduce along the arg-max axis, with parallel operations along
2039     iteratorTypes[axis] = utils::IteratorType::reduction;
2045       if (axis != i)
2063               rewriter.create<linalg::IndexOp>(loc, axis));