diff --git a/more_math/Parser/UnifiedMathVisitor.py b/more_math/Parser/UnifiedMathVisitor.py index c0ff266..cf2203a 100644 --- a/more_math/Parser/UnifiedMathVisitor.py +++ b/more_math/Parser/UnifiedMathVisitor.py @@ -2772,3 +2772,19 @@ class UnifiedMathVisitor(MathExprVisitor): return noise + def visitPadFunc(self,ctx): + val = self._promote_to_tensor((yield ctx.expr(0))) + pad_val = yield ctx.expr(1) + if self._is_tensor(pad_val): + pad = [int(x) for x in pad_val.flatten().tolist()] + elif self._is_list(pad_val): + pad = [int(x) for x in pad_val] + else: + raise ValueError(f"{ctx.start.line}:{ctx.start.column}: Pad amount must be a list or tensor of integers") + if len(pad) % 2 != 0: + raise ValueError(f"{ctx.start.line}:{ctx.start.column}: Pad amount list must have an even number of elements") + + reversed_pad = [] + for i in range(len(pad) - 1, 0, -2): + reversed_pad.extend([pad[i-1], pad[i]]) + return F.pad(val, reversed_pad)