Fix type infer of Expand

kendryte · Mar 6, 2025 · 0273c62 · 0273c62
1 parent 73dcee7
commit 0273c62
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 3 deletions.
diff --git a/src/Nncase.Core/Evaluator/TypeInference.cs b/src/Nncase.Core/Evaluator/TypeInference.cs
@@ -590,4 +590,26 @@ public static IRType[] BroadcastDistributeTypes(params IRType[] types)
 
         return types;
     }
+
+    public static Shape ExpandShape(Shape inShape, Shape expandShape)
+    {
+        if (inShape.IsUnranked || expandShape.IsUnranked)
+        {
+            return Shape.Unranked;
+        }
+
+        var dimExtends = expandShape.Rank - inShape.Rank;
+        var newDims = expandShape.ToArray();
+
+        // dimsExtends may be negative
+        for (int i = Math.Max(0, dimExtends); i < newDims.Length; i++)
+        {
+            var inDimIndex = i - dimExtends;
+            ref var dimValue = ref newDims[i];
+            dimValue = Dimension.Select(dimValue, 1L, inShape[inDimIndex], dimValue);
+        }
+
+        newDims = inShape.Take(dimExtends < 0 ? -dimExtends : 0).Concat(newDims).ToArray();
+        return new Shape(newDims);
+    }
 }
diff --git a/src/Nncase.Evaluator/Tensors/Expand.cs b/src/Nncase.Evaluator/Tensors/Expand.cs
@@ -62,15 +62,17 @@ public IRType Visit(ITypeInferenceContext context, Expand target)
     private IRType Visit(ITypeInferenceContext context, Expand target, TensorType input, TensorType shape)
     {
         var shapeExpr = context.GetArgument(target, Expand.Shape);
-        return input with { Shape = Shape.FromExpr(shapeExpr) };
+        var newShape = TypeInference.ExpandShape(input.Shape, Shape.FromExpr(shapeExpr));
+        return input with { Shape = newShape };
     }
 
     private IRType Visit(ITypeInferenceContext context, Expand target, DistributedType input, TensorType shape)
     {
         var invalid = new InvalidType(input.ToString());
-        var newShape = Shape.FromExpr(context.GetArgument(target, Expand.Shape));
-        if (newShape.IsRanked)
+        var shapeExpr = Shape.FromExpr(context.GetArgument(target, Expand.Shape));
+        if (input.TensorType.Shape.IsRanked && shapeExpr.IsRanked)
         {
+            var newShape = TypeInference.ExpandShape(input.TensorType.Shape, shapeExpr);
             var ndsbp = new SBP[input.Placement.Rank];
             for (int i = 0; i < input.Placement.Rank; i++)
             {