From 412ea5e89be8ee1ac6b5601651b9c8126b666fe2 Mon Sep 17 00:00:00 2001 From: mcDandy Date: Tue, 17 Feb 2026 00:34:51 +0100 Subject: [PATCH] shape returns list now --- more_math/Parser/UnifiedMathVisitor.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/more_math/Parser/UnifiedMathVisitor.py b/more_math/Parser/UnifiedMathVisitor.py index e2d56c9..f73808f 100644 --- a/more_math/Parser/UnifiedMathVisitor.py +++ b/more_math/Parser/UnifiedMathVisitor.py @@ -2376,11 +2376,14 @@ class UnifiedMathVisitor(MathExprVisitor): def visitShapeFunc(self, ctx): val = (yield ctx.expr()) if self._is_tensor(val): - return torch.tensor(list(val.shape), dtype=torch.long, device=self.device) + # Return shape as a 1D tensor of integers + return list(val.shape) elif self._is_list(val): - return torch.tensor([len(val)], dtype=torch.long, device=self.device) + # Return list length as a single-element tensor + return [len(val)] else: - return torch.tensor([], dtype=torch.long, device=self.device) + # Scalar has shape [] + return [] def _bitwise_op(self, a, b, torch_op, scalar_op): """Binary bitwise operation handler supporting tensors, lists, and scalars."""