From d74304d0f24c73d0af39dff2cbb3414dad5101fd Mon Sep 17 00:00:00 2001 From: mcDandy Date: Sat, 28 Feb 2026 17:44:37 +0100 Subject: [PATCH] fix shape --- more_math/Parser/UnifiedMathVisitor.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/more_math/Parser/UnifiedMathVisitor.py b/more_math/Parser/UnifiedMathVisitor.py index 549bfbd..b363efd 100644 --- a/more_math/Parser/UnifiedMathVisitor.py +++ b/more_math/Parser/UnifiedMathVisitor.py @@ -2382,7 +2382,12 @@ class UnifiedMathVisitor(MathExprVisitor): type = (yield ctx.expr(1)).dtype if ctx.expr(1) else None shape_val = yield ctx.indexExpr() - shape = [self._to_int(shape_val, ctx, "tensor")] + if self._is_list(shape_val): + shape = [self._to_int(v, ctx, "tensor") for v in shape_val] + elif self._is_tensor(shape_val): + shape = [self._to_int(v, ctx, "tensor") for v in shape_val.flatten().tolist()] + else: + shape = [self._to_int(shape_val, ctx, "tensor")] return torch.full(shape, value, device=self.device,dtype=type)