diff --git a/more_math/Parser/MathExpr.g4 b/more_math/Parser/MathExpr.g4 index f05c071..e1b5470 100644 --- a/more_math/Parser/MathExpr.g4 +++ b/more_math/Parser/MathExpr.g4 @@ -167,7 +167,9 @@ func1: | DILATE LPAREN expr (COMMA expr)? RPAREN # DilateFunc | ERODE LPAREN expr (COMMA expr)? RPAREN # ErodeFunc | MORPH_OPEN LPAREN expr (COMMA expr)? RPAREN # MorphOpenFunc - | MORPH_CLOSE LPAREN expr (COMMA expr)? RPAREN # MorphCloseFunc; + | MORPH_CLOSE LPAREN expr (COMMA expr)? RPAREN # MorphCloseFunc + | INT LPAREN expr RPAREN # IntFunc; + | FLOAT LPAREN expr RPAREN # FloatFunc; func2: POWE LPAREN expr COMMA expr RPAREN # PowFunc @@ -371,6 +373,8 @@ CLEAR: 'stack_clear'; HAS: 'stack_has'; GET: 'stack_get'; CONCAT: 'cat'|'concat'|'concatenate'; +INT: 'int'; +FLOAT: 'float'; UPPER: 'upper'; LOWER: 'lower'; diff --git a/more_math/Parser/UnifiedMathVisitor.py b/more_math/Parser/UnifiedMathVisitor.py index 2808cb0..ffb26df 100644 --- a/more_math/Parser/UnifiedMathVisitor.py +++ b/more_math/Parser/UnifiedMathVisitor.py @@ -3399,5 +3399,38 @@ class UnifiedMathVisitor(MathExprVisitor): d = int(dim_val.item()) if self._is_tensor(dim_val) else int(dim_val) return torch.cat(tensors, dim=d) + def visitIntFunc(self, ctx): + val = yield ctx.expr() + def as_int(ctx,val): + if self._is_tensor(val): + return val.to(torch.int32).contiguous() + + if self._is_list(val): + return [as_int(ctx,x) for x in val] + + if isinstance(val, str): + return int(float(val)) + + if val is None: + raise ValueError(f"{ctx.start.line}:{ctx.start.column}: Cannot convert None to a number") + + return int(float(val)) + return as_int(ctx,val) - + def visitFloatFunc(self, ctx): + val = yield ctx.expr() + def as_float(ctx,val): + if self._is_tensor(val): + return val.to(torch.float).contiguous() + + if self._is_list(val): + return [as_float(ctx,x) for x in val] + + if isinstance(val, str): + return float(val) + + if val is None: + raise ValueError(f"{ctx.start.line}:{ctx.start.column}: Cannot convert None to a number") + + return float(val) + return as_float(ctx,val) \ No newline at end of file