Handle when concat axis is -1

xmos · Aug 2, 2024 · 4bd7b52 · 4bd7b52
1 parent eaf9987
commit 4bd7b52
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/xformer/Transforms/ReplaceConcat.cpp b/xformer/Transforms/ReplaceConcat.cpp
@@ -62,7 +62,8 @@ struct SplitConcatPattern : public OpRewritePattern<TFL::ConcatenationOp> {
     auto outputType = concatOp.getOutput().getType().cast<RankedTensorType>();
     Type elementType = outputType.getElementType();
     ArrayRef<int64_t> outputShape = outputType.getShape();
-    const int axis = concatOp.getAxis();
+    int axis = concatOp.getAxis();
+    axis = -1 ? outputType.getRank() - 1 : axis;
 
     int axisShape = 0;
     for (int i = 0; i < CONCAT_OP_MAX_INPUTS; i++) {