jaclang 0.8.7__py3-none-any.whl → 0.8.8__py3-none-any.whl

jaclang/compiler/passes/main/pyast_gen_pass.py

@@ -1000,6 +1000,9 @@ class PyastGenPass(UniPass):
         pass
 
     def exit_func_signature(self, node: uni.FuncSignature) -> None:
+        posonlyargs = [i.gen.py_ast[0] for i in node.posonly_params]
+        vararg = node.varargs.gen.py_ast[0] if node.varargs else None
+        kwarg = node.kwargs.gen.py_ast[0] if node.kwargs else None
         params = (
             [self.sync(ast3.arg(arg="self", annotation=None))]
             if (abl := node.parent)
@@ -1009,29 +1012,33 @@ class PyastGenPass(UniPass):
             and not node.is_in_py_class
             else []
         )
-        vararg = None
-        kwarg = None
-        for i in node.params:
-            if i.unpack and i.unpack.value == "*":
-                vararg = i.gen.py_ast[0]
-            elif i.unpack and i.unpack.value == "**":
-                kwarg = i.gen.py_ast[0]
+        if posonlyargs:
+            posonlyargs = params + posonlyargs
+            params = [i.gen.py_ast[0] for i in node.params]
+        else:
+            params = params + [i.gen.py_ast[0] for i in node.params]
+        defaults = []
+        for i in [*node.posonly_params, *node.params]:
+            if i.value:
+                defaults.append(cast(ast3.expr, i.value.gen.py_ast[0]))
+        kwonly_args = [i.gen.py_ast[0] for i in node.kwonlyargs]
+        # kw_defaults must be the same length as kwonlyargs
+        # it will have None for args that don't have defaults
+        kw_defaults: list[ast3.expr | None] = []
+        for i in node.kwonlyargs:
+            if i.value:
+                kw_defaults.append(cast(ast3.expr, i.value.gen.py_ast[0]))
             else:
-                (
-                    params.append(i.gen.py_ast[0])
-                    if isinstance(i.gen.py_ast[0], ast3.arg)
-                    else self.ice("This list should only be Args")
-                )
-        defaults = [x.value.gen.py_ast[0] for x in node.params if x.value]
+                kw_defaults.append(None)
         node.gen.py_ast = [
             self.sync(
                 ast3.arguments(
-                    posonlyargs=[],
+                    posonlyargs=[cast(ast3.arg, param) for param in posonlyargs],
                     args=[cast(ast3.arg, param) for param in params],
-                    kwonlyargs=[],
+                    kwonlyargs=kwonly_args,
                     vararg=cast(ast3.arg, vararg) if vararg else None,
                     kwarg=cast(ast3.arg, kwarg) if kwarg else None,
-                    kw_defaults=[],
+                    kw_defaults=kw_defaults,
                     defaults=[cast(ast3.expr, default) for default in defaults],
                 )
             )
@@ -1467,20 +1474,23 @@ class PyastGenPass(UniPass):
         ]
 
     def exit_assert_stmt(self, node: uni.AssertStmt) -> None:
-        node.gen.py_ast = [
-            self.sync(
-                ast3.Assert(
-                    test=cast(ast3.expr, node.condition.gen.py_ast[0]),
-                    msg=(
-                        cast(ast3.expr, node.error_msg.gen.py_ast[0])
-                        if node.error_msg
-                        else None
-                    ),
+        if isinstance(node.parent, uni.Test):
+            self.assert_helper(node)
+        else:
+            node.gen.py_ast = [
+                self.sync(
+                    ast3.Assert(
+                        test=cast(ast3.expr, node.condition.gen.py_ast[0]),
+                        msg=(
+                            cast(ast3.expr, node.error_msg.gen.py_ast[0])
+                            if node.error_msg
+                            else None
+                        ),
+                    )
                 )
-            )
-        ]
+            ]
 
-    def exit_check_stmt(self, node: uni.CheckStmt) -> None:
+    def assert_helper(self, node: uni.AssertStmt) -> None:
         """Sub objects.
 
         target: ExprType,
@@ -1536,16 +1546,16 @@ class PyastGenPass(UniPass):
 
         # By default the check expression will become assertTrue(<expr>), unless any pattern detected.
         assert_func_name = "assertTrue"
-        assert_args_list = node.target.gen.py_ast
+        assert_args_list = node.condition.gen.py_ast
 
         # Compare operations. Note that We're only considering the compare
         # operation with a single operation ie. a < b < c is  ignored here.
         if (
-            isinstance(node.target, uni.CompareExpr)
-            and isinstance(node.target.gen.py_ast[0], ast3.Compare)
-            and len(node.target.ops) == 1
+            isinstance(node.condition, uni.CompareExpr)
+            and isinstance(node.condition.gen.py_ast[0], ast3.Compare)
+            and len(node.condition.ops) == 1
         ):
-            expr: uni.CompareExpr = node.target
+            expr: uni.CompareExpr = node.condition
             opty: uni.Token = expr.ops[0]
 
             optype2fn = {
@@ -1579,10 +1589,10 @@ class PyastGenPass(UniPass):
                     assert_args_list.pop()
 
         # Check if 'isinstance' is called.
-        elif isinstance(node.target, uni.FuncCall) and isinstance(
-            node.target.gen.py_ast[0], ast3.Call
+        elif isinstance(node.condition, uni.FuncCall) and isinstance(
+            node.condition.gen.py_ast[0], ast3.Call
         ):
-            res = check_node_isinstance_call(node.target)
+            res = check_node_isinstance_call(node.condition)
             if res.isit:
                 # These assertions will make mypy happy.
                 assert isinstance(res.inst, ast3.AST)
@@ -1592,12 +1602,12 @@ class PyastGenPass(UniPass):
 
         # Check if 'not isinstance(<expr>, <expr>)' is called.
         elif (
-            isinstance(node.target, uni.UnaryExpr)
-            and isinstance(node.target, uni.UnaryExpr)
-            and isinstance(node.target.operand, uni.FuncCall)
-            and isinstance(node.target.operand, uni.UnaryExpr)
+            isinstance(node.condition, uni.UnaryExpr)
+            and isinstance(node.condition, uni.UnaryExpr)
+            and isinstance(node.condition.operand, uni.FuncCall)
+            and isinstance(node.condition.operand, uni.UnaryExpr)
         ):
-            res = check_node_isinstance_call(node.target.operand)
+            res = check_node_isinstance_call(node.condition.operand)
             if res.isit:
                 # These assertions will make mypy happy.
                 assert isinstance(res.inst, ast3.AST)
@@ -1610,14 +1620,14 @@ class PyastGenPass(UniPass):
         # the almost equal functionality (snice there is no almost equal operator in jac and never needed ig.).
 
         # Check if 'almostEqual' is called.
-        if isinstance(node.target, uni.FuncCall) and isinstance(
-            node.target.gen.py_ast[0], ast3.Call
+        if isinstance(node.condition, uni.FuncCall) and isinstance(
+            node.condition.gen.py_ast[0], ast3.Call
         ):
-            func = node.target.target
+            func = node.condition.target
             if isinstance(func, uni.Name) and func.value == "almostEqual":
                 assert_func_name = "assertAlmostEqual"
                 assert_args_list = []
-                for param in node.target.params:
+                for param in node.condition.params:
                     assert_args_list.append(param.gen.py_ast[0])
 
         # assert_func_expr = "Con.JAC_CHECK.value.assertXXX"
@@ -2718,14 +2728,22 @@ class PyastGenPass(UniPass):
                     keywords=[],
                 )
             )
-
-        pynode = self.sync(
-            ast3.Call(
-                func=self.jaclib_obj("refs"),
-                args=[pynode],
-                keywords=[],
+        if node.is_async:
+            pynode = self.sync(
+                ast3.Call(
+                    func=self.jaclib_obj("arefs"),
+                    args=[pynode],
+                    keywords=[],
+                )
+            )
+        else:
+            pynode = self.sync(
+                ast3.Call(
+                    func=self.jaclib_obj("refs"),
+                    args=[pynode],
+                    keywords=[],
+                )
             )
-        )
 
         node.gen.py_ast = [pynode]
 

jaclang/compiler/passes/main/pyast_load_pass.py

@@ -16,7 +16,7 @@ from __future__ import annotations
 
 import ast as py_ast
 import os
-from typing import Optional, Sequence, TYPE_CHECKING, TypeAlias, TypeVar
+from typing import Optional, Sequence, TYPE_CHECKING, TypeAlias, TypeVar, cast
 
 import jaclang.compiler.unitree as uni
 from jaclang.compiler.constant import Tokens as Tok
@@ -2091,6 +2091,7 @@ class PyastBuildPass(Transform[uni.PythonModuleAst, uni.Module]):
         """Process python node.
 
         class arguments(AST):
+            posonlyargs: list[arg]
             args: list[arg]
             vararg: arg | None
             kwonlyargs: list[arg]
@@ -2098,9 +2099,23 @@ class PyastBuildPass(Transform[uni.PythonModuleAst, uni.Module]):
             kwarg: arg | None
             defaults: list[expr]
         """
-        args = [self.convert(arg) for arg in node.args]
+
+        def _apply_kind(params: list, kind: uni.ParamKind) -> list:
+            for param in params:
+                cast(uni.ParamVar, param).param_kind = kind
+            return params
+
+        posonlyargs = _apply_kind(
+            [self.convert(arg) for arg in node.posonlyargs], uni.ParamKind.POSONLY
+        )
+        args = _apply_kind(
+            [self.convert(arg) for arg in node.args], uni.ParamKind.NORMAL
+        )
+
         vararg = self.convert(node.vararg) if node.vararg else None
+
         if vararg and isinstance(vararg, uni.ParamVar):
+            vararg.param_kind = uni.ParamKind.VARARG
             vararg.unpack = uni.Token(
                 orig_src=self.orig_src,
                 name=Tok.STAR_MUL,
@@ -2113,7 +2128,10 @@ class PyastBuildPass(Transform[uni.PythonModuleAst, uni.Module]):
                 pos_end=0,
             )
             vararg.add_kids_left([vararg.unpack])
-        kwonlyargs = [self.convert(arg) for arg in node.kwonlyargs]
+
+        kwonlyargs = _apply_kind(
+            [self.convert(arg) for arg in node.kwonlyargs], uni.ParamKind.KWONLY
+        )
         for i in range(len(kwonlyargs)):
             kwa = kwonlyargs[i]
             kwd = node.kw_defaults[i]
@@ -2127,6 +2145,7 @@ class PyastBuildPass(Transform[uni.PythonModuleAst, uni.Module]):
                 kwa.add_kids_right([kwa.value])
         kwarg = self.convert(node.kwarg) if node.kwarg else None
         if kwarg and isinstance(kwarg, uni.ParamVar):
+            kwarg.param_kind = uni.ParamKind.KWARG
             kwarg.unpack = uni.Token(
                 orig_src=self.orig_src,
                 name=Tok.STAR_POW,
@@ -2140,29 +2159,42 @@ class PyastBuildPass(Transform[uni.PythonModuleAst, uni.Module]):
             )
             kwarg.add_kids_left([kwarg.unpack])
         defaults = [self.convert(expr) for expr in node.defaults]
-        params = [*args]
-        for param, default in zip(params[::-1], defaults[::-1]):
-            if isinstance(default, uni.Expr) and isinstance(param, uni.ParamVar):
-                param.value = default
-                param.add_kids_right([default])
-        if vararg:
-            params.append(vararg)
-        params += kwonlyargs
-        if kwarg:
-            params.append(kwarg)
-        params += defaults
-
-        valid_params = [param for param in params if isinstance(param, uni.ParamVar)]
-        if valid_params:
-            fs_params = valid_params
+        # iterate reverse to match from the end
+        for para in [*posonlyargs, *args][::-1]:
+            if not defaults:
+                break
+            default = defaults.pop()
+            if (
+                default
+                and isinstance(para, uni.ParamVar)
+                and isinstance(default, uni.Expr)
+            ):
+                para.value = default
+                para.add_kids_right([para.value])
+
+        if kwonlyargs or args or posonlyargs or vararg or kwarg:
+            kids = []
+            kids.extend(posonlyargs) if posonlyargs else None
+            kids.extend(args) if args else None
+            kids.append(vararg) if vararg else None
+            kids.extend(kwonlyargs) if kwonlyargs else None
+            kids.append(kwarg) if kwarg else None
             return uni.FuncSignature(
-                params=fs_params,
+                posonly_params=posonlyargs,
+                params=args,
+                varargs=vararg,
+                kwonlyargs=kwonlyargs,
+                kwargs=kwarg,
                 return_type=None,
-                kid=fs_params,
+                kid=kids,
             )
         else:
             return uni.FuncSignature(
+                posonly_params=posonlyargs,
                 params=[],
+                varargs=vararg,
+                kwonlyargs=kwonlyargs,
+                kwargs=kwarg,
                 return_type=None,
                 kid=[self.operator(Tok.LPAREN, "("), self.operator(Tok.RPAREN, ")")],
             )

jaclang/compiler/passes/main/sym_tab_build_pass.py

@@ -75,7 +75,7 @@ class SymTabBuildPass(UniPass):
     def exit_module_path(self, node: uni.ModulePath) -> None:
         if node.alias:
             node.alias.sym_tab.def_insert(node.alias, single_decl="import")
-        elif node.path and isinstance(node.path[0], uni.Name):
+        elif node.path:
             if node.parent_of_type(uni.Import) and not (
                 node.parent_of_type(uni.Import).from_loc
                 and node.parent_of_type(uni.Import).is_jac

jaclang/compiler/passes/main/tests/fixtures/checker/import_sym.jac

@@ -0,0 +1,2 @@
+
+import from imported_sym { foo }

jaclang/compiler/passes/main/tests/fixtures/checker/import_sym_test.jac

@@ -0,0 +1,6 @@
+import from import_sym { foo }
+
+with entry {
+    a: str = foo();  # <-- Ok
+    b: int = foo();  # <-- Error
+}

jaclang/compiler/passes/main/tests/fixtures/checker/imported_sym.jac

@@ -0,0 +1,5 @@
+
+
+def foo() -> str {
+    return "foo";
+}

jaclang/compiler/passes/main/tests/fixtures/checker_arg_param_match.jac

@@ -0,0 +1,37 @@
+
+class Foo {
+    def bar(self: Foo, a:int) -> None {}
+    static def baz(self: int, a: int) -> None {}
+}
+
+def foo(a: int, b: int, /, c:int, d:int, *args:int, e:int, f:int, **kwargs:int) -> None {}
+def bar(a: int, b: int, /, c:int, d:int, e:int, f:int) -> None {}
+def baz(a: int, *, b:int) -> None {}
+
+with entry {
+    f = Foo();
+
+    f.bar();
+    f.bar(1);
+    f.bar(1, 2);
+
+    f.baz();
+    f.baz(1);
+    f.baz(1, 2);
+
+    foo(1, 2, 3, d=4, e=5, f=6, g=7, h=8); # c is positional and d is named
+    foo(1, 2, 3, 4, 5, 6, 7, 8, e=5, f=6); # marching extra with *args
+    foo(1, 2, d=3, e=4, f=5, c=4);  # order does not matter for named
+
+    foo(1, 2, 3, d=4, e=5, g=7, h=8); # missing argument 'f'
+    foo(1, b=2, c=3, d=4, e=5, f=6); # b is positional only
+
+    bar(1, 2, 3, 4, 5, f=6);
+    bar(1, 2, 3, 4, 5, 6, 7, 8, 9); # too many args
+    bar(1, 2, 3, 4, 5, 6, c=3); # already matched
+    bar(1, 2, 3, 4, 5, 6, h=1); # h is not matched
+
+    baz(a=1, b=2);
+    baz(1, b=2); # a can be both positional and keyword
+    baz(1, 2);  # 'b' can only be keyword arg
+}

jaclang/compiler/passes/main/tests/fixtures/checker_arity.jac

@@ -0,0 +1,18 @@
+
+class Foo {
+    def first_is_self(self: Foo) -> None {}
+
+    def with_default_args(self: Foo, a:int, b:int=42) -> None {}
+}
+
+
+with entry {
+    f = Foo();
+    f.first_is_self();  # <-- Ok
+    f.first_is_self(f); # <-- Error
+
+    f.with_default_args(1);       # <-- Ok
+    f.with_default_args(1, 2);    # <-- Ok
+    f.with_default_args(1, 2, 3); # <-- Error
+    f.with_default_args();        # <-- Error
+}

jaclang/compiler/passes/main/tests/fixtures/checker_cat_is_animal.jac

@@ -0,0 +1,18 @@
+
+obj Animal {}
+obj Cat(Animal) {}
+obj Lion(Cat) {}
+
+obj NotAnimal {}
+
+def animal_func(a: Animal) -> None {}
+
+with entry {
+    cat: Cat = Cat();
+    lion: Lion = Lion();
+    not_animal: NotAnimal = NotAnimal();
+
+    animal_func(cat);        # <-- Ok
+    animal_func(lion);       # <-- Ok
+    animal_func(not_animal); # <-- Error
+}

jaclang/compiler/passes/main/tests/fixtures/checker_float.jac

@@ -0,0 +1,7 @@
+
+
+with entry {
+    pi = 3.14;     # <-- infer the type
+    f: float = pi; # <-- OK
+    s: str = pi;   # <-- Error
+}

jaclang/compiler/passes/main/tests/fixtures/checker_param_types.jac

@@ -0,0 +1,11 @@
+
+node A {}
+node B {}
+
+def foo(a: A) -> None { }
+
+
+with entry {
+    foo(A()); # <-- Ok
+    foo(B()); # <-- Error
+}

jaclang/compiler/passes/main/tests/fixtures/checker_self_type.jac

@@ -0,0 +1,9 @@
+
+node Foo {
+    has i: int = 0;
+
+    def foo() {
+     y: int = self.i; # <-- Ok
+     x: str = self.i; # <-- Error
+    }
+}

jaclang/compiler/passes/main/tests/fixtures/checker_sym_inherit.jac

@@ -0,0 +1,42 @@
+
+# -----------------------------------------------------------------------------
+# Simple Inheritance
+# -----------------------------------------------------------------------------
+
+node Parent {
+    has val: int = 42;
+}
+
+node Child (Parent) {
+
+}
+
+
+with entry {
+    c = Child();
+    c.val = 42;     # <-- Ok
+    c.val = "str";  # <-- Error
+}
+
+
+# -----------------------------------------------------------------------------
+# A Complex Inheritance
+# -----------------------------------------------------------------------------
+
+node Animal {
+    has name: str;
+}
+
+node Cat(Animal) {
+}
+
+node Lion(Cat) {
+}
+
+
+with entry {
+    l = Lion();
+
+    l.name = "Simba";  # <-- Ok
+    l.name = 42;       # <-- Error
+}

jaclang/compiler/passes/main/tests/fixtures/predynamo_fix3.jac

@@ -0,0 +1,43 @@
+import torch;
+class Cfg {
+    def __init__(self: Cfg, max_position_embeddings: Any = 8, has_original: Any = False) {
+        self.max_position_embeddings = max_position_embeddings;
+        if has_original {
+            self.original_max_position_embeddings = (max_position_embeddings // 2);
+        }
+    }
+}
+class ToyModel(torch.nn.Module) {
+    def __init__(self: ToyModel, cfg: Cfg) {
+        super.init();
+        self.config = cfg;
+        self.long_inv_freq = torch.tensor([10.0, 20.0, 30.0]);
+        self.original_inv_freq = torch.tensor([1.0, 2.0, 3.0]);
+    }
+    def rope_init_fn(self: ToyModel, cfg: Any, device: Any, seq_len: int) {
+        return (torch.arange(1, 4, dtype=torch.float32, device=device), None);
+    }
+    def _longrope_frequency_update(
+        self: ToyModel,
+        position_ids: Any,
+        device: Any = 'cpu'
+    ) {
+        seq_len = (torch.max(position_ids) + 1);
+        if hasattr(self.config, 'original_max_position_embeddings') {
+            original_max_position_embeddings = self.config.original_max_position_embeddings;
+        } else {
+            original_max_position_embeddings = self.config.max_position_embeddings;
+        }
+        if (seq_len > original_max_position_embeddings) {
+            self.register_buffer('inv_freq', self.long_inv_freq, persistent=False);
+        } else {
+            self.register_buffer(
+                'inv_freq',
+                self.original_inv_freq.to(device),
+                persistent=False
+            );
+        }
+        return self.inv_freq;
+    }
+}
+

jaclang/compiler/passes/main/tests/fixtures/predynamo_where_assign.jac

@@ -0,0 +1,13 @@
+import torch;
+
+def with_breaks(a: Any, b: Any) {
+    x = (a / (torch.abs(a) + 1));
+    b = (b * 1.414);
+    if (b.sum() < 0) {
+        b = (b * -1);
+    } else{
+            b = (b * -2);
+    }
+    
+    return (x * b);
+}

jaclang/compiler/passes/main/tests/fixtures/predynamo_where_return.jac

@@ -0,0 +1,11 @@
+import torch;
+
+def with_breaks(a: Any, b: Any) {
+    x = (a / (torch.abs(a) + 1));
+    b = (b * 1.414);
+    if (b.sum() < 0) {
+        return x * (b * -1);
+    } else{
+        return x * (b * -2);
+    }
+}