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

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
@@ -96,6 +96,8 @@ class PyastBuildPass(Transform[uni.PythonModuleAst, uni.Module]):
             body: list[stmt]
             type_ignores: list[TypeIgnore]
         """
+        if not node.body:
+            return uni.Module.make_stub(inject_src=self.ir_in)
         elements: list[uni.UniNode] = [self.convert(i) for i in node.body]
         elements[0] = (
             elements[0].expr
@@ -2091,6 +2093,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 +2101,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 +2130,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 +2147,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 +2161,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);
+    }
+}

jaclang/compiler/passes/main/tests/test_checker_pass.py

@@ -28,6 +28,18 @@ class TypeCheckerPassTests(TestCase):
                  ^^^^^^^^^^^^^^^^^^^^^^
         """, program.errors_had[1].pretty_print())
 
+    def test_float_types(self) -> None:
+        program = JacProgram()
+        mod = program.compile(self.fixture_abs_path("checker_float.jac"))
+        TypeCheckPass(ir_in=mod, prog=program)
+        self.assertEqual(len(program.errors_had), 1)
+        self._assert_error_pretty_found("""
+            f: float = pi; # <-- OK
+            s: str = pi;   # <-- Error
+            ^^^^^^^^^^^
+        """, program.errors_had[0].pretty_print())
+
+
     def test_infer_type_of_assignment(self) -> None:
         program = JacProgram()
         mod = program.compile(self.fixture_abs_path("infer_type_assignment.jac"))
@@ -49,6 +61,17 @@ class TypeCheckerPassTests(TestCase):
           ^^^^^^^^^^^^^^^^^^^
         """, program.errors_had[0].pretty_print())
 
+    def test_imported_sym(self) -> None:
+        program = JacProgram()
+        mod = program.compile(self.fixture_abs_path("checker/import_sym_test.jac"))
+        TypeCheckPass(ir_in=mod, prog=program)
+        self.assertEqual(len(program.errors_had), 1)
+        self._assert_error_pretty_found("""
+          a: str = foo();  # <-- Ok
+          b: int = foo();  # <-- Error
+          ^^^^^^^^^^^^^^
+        """, program.errors_had[0].pretty_print())
+
     def test_member_access_type_infered(self) -> None:
         program = JacProgram()
         mod = program.compile(self.fixture_abs_path("member_access_type_inferred.jac"))
@@ -59,6 +82,22 @@ class TypeCheckerPassTests(TestCase):
           ^^^^^^^^^
         """, program.errors_had[0].pretty_print())
 
+    def test_inherited_symbol(self) -> None:
+        program = JacProgram()
+        mod = program.compile(self.fixture_abs_path("checker_sym_inherit.jac"))
+        TypeCheckPass(ir_in=mod, prog=program)
+        self.assertEqual(len(program.errors_had), 2)
+        self._assert_error_pretty_found("""
+          c.val = 42;     # <-- Ok
+          c.val = "str";  # <-- Error
+          ^^^^^^^^^^^^^
+        """, program.errors_had[0].pretty_print())
+        self._assert_error_pretty_found("""
+          l.name = "Simba";  # <-- Ok
+          l.name = 42;       # <-- Error
+          ^^^^^^^^^^^
+        """, program.errors_had[1].pretty_print())
+
     def test_import_symbol_type_infer(self) -> None:
         program = JacProgram()
         mod = program.compile(self.fixture_abs_path("import_symbol_type_infer.jac"))
@@ -104,6 +143,144 @@ class TypeCheckerPassTests(TestCase):
             ^^^^^^^^^^^^^^^^
         """, program.errors_had[0].pretty_print())
 
+    def test_arity(self) -> None:
+        path = self.fixture_abs_path("checker_arity.jac")
+        program = JacProgram()
+        mod = program.compile(path)
+        TypeCheckPass(ir_in=mod, prog=program)
+        self.assertEqual(len(program.errors_had), 3)
+        self._assert_error_pretty_found("""
+            f.first_is_self(f); # <-- Error
+                            ^
+        """, program.errors_had[0].pretty_print())
+        self._assert_error_pretty_found("""
+            f.with_default_args(1, 2, 3); # <-- Error
+                                      ^
+        """, program.errors_had[1].pretty_print())
+        self._assert_error_pretty_found("""
+            f.with_default_args();        # <-- Error
+            ^^^^^^^^^^^^^^^^^^^^^
+        """, program.errors_had[2].pretty_print())
+
+    def test_param_types(self) -> None:
+        path = self.fixture_abs_path("checker_param_types.jac")
+        program = JacProgram()
+        mod = program.compile(path)
+        TypeCheckPass(ir_in=mod, prog=program)
+        self.assertEqual(len(program.errors_had), 1)
+        self._assert_error_pretty_found("""
+            foo(A()); # <-- Ok
+            foo(B()); # <-- Error
+                ^^^
+        """, program.errors_had[0].pretty_print())
+
+    def test_param_arg_match(self) -> None:
+        program = JacProgram()
+        path = self.fixture_abs_path("checker_arg_param_match.jac")
+        mod = program.compile(path)
+        TypeCheckPass(ir_in=mod, prog=program)
+        self.assertEqual(len(program.errors_had), 13)
+
+        expected_errors = [
+            """
+            Not all required parameters were provided in the function call: 'a'
+                     f = Foo();
+                     f.bar();
+                     ^^^^^^^
+            """,
+            """
+            Too many positional arguments
+                     f.bar();
+                     f.bar(1);
+                     f.bar(1, 2);
+                              ^
+            """,
+            """
+            Not all required parameters were provided in the function call: 'self', 'a'
+                     f.bar(1, 2);
+                     f.baz();
+                     ^^^^^^^
+            """,
+            """
+            Not all required parameters were provided in the function call: 'a'
+                     f.baz();
+                     f.baz(1);
+                     ^^^^^^^^
+            """,
+            """
+            Not all required parameters were provided in the function call: 'f'
+                     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'
+                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+            """,
+            """
+            Positional only parameter 'b' cannot be matched with a named argument
+                     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
+                            ^^^
+            """,
+            """
+            Too many positional arguments
+                     bar(1, 2, 3, 4, 5, f=6);
+                     bar(1, 2, 3, 4, 5, 6, 7, 8, 9); # too many args
+                                           ^
+            """,
+            """
+            Too many positional arguments
+                     bar(1, 2, 3, 4, 5, f=6);
+                     bar(1, 2, 3, 4, 5, 6, 7, 8, 9); # too many args
+                                              ^
+            """,
+            """
+            Too many positional arguments
+                     bar(1, 2, 3, 4, 5, f=6);
+                     bar(1, 2, 3, 4, 5, 6, 7, 8, 9); # too many args
+                                                 ^
+            """,
+            """
+            Parameter 'c' already matched
+                     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
+                                           ^^^
+            """,
+            """
+            Named argument 'h' does not match any parameter
+                     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
+                                           ^^^
+            """,
+            """
+            Too many positional arguments
+                     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
+                            ^
+            """,
+            """
+            Not all required parameters were provided in the function call: 'b'
+                     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
+                     ^^^^^^^^^
+            """,
+        ]
+
+        for i, expected in enumerate(expected_errors):
+            self._assert_error_pretty_found(expected, program.errors_had[i].pretty_print())
+
+    def test_self_type_inference(self) -> None:
+        path = self.fixture_abs_path("checker_self_type.jac")
+        program = JacProgram()
+        mod = program.compile(path)
+        TypeCheckPass(ir_in=mod, prog=program)
+        self.assertEqual(len(program.errors_had), 1)
+        self._assert_error_pretty_found("""
+          x: str = self.i; # <-- Error
+          ^^^^^^^^^^^^^^^
+        """, program.errors_had[0].pretty_print())
+
     def test_binary_op(self) -> None:
         program = JacProgram()
         mod = program.compile(self.fixture_abs_path("checker_binary_op.jac"))
@@ -140,6 +317,19 @@ class TypeCheckerPassTests(TestCase):
             ^^^^^^^^^^^^^^
         """, program.errors_had[0].pretty_print())
 
+    def test_checker_cat_is_animal(self) -> None:
+        path = self.fixture_abs_path("checker_cat_is_animal.jac")
+        program = JacProgram()
+        mod = program.compile(path)
+        TypeCheckPass(ir_in=mod, prog=program)
+        self.assertEqual(len(program.errors_had), 1)
+        self._assert_error_pretty_found("""
+            animal_func(cat);        # <-- Ok
+            animal_func(lion);       # <-- Ok
+            animal_func(not_animal); # <-- Error
+                        ^^^^^^^^^^
+        """, program.errors_had[0].pretty_print())
+
     def test_checker_import_missing_module(self) -> None:
         path = self.fixture_abs_path("checker_import_missing_module.jac")
         program = JacProgram()

jaclang/compiler/passes/main/tests/test_predynamo_pass.py

@@ -0,0 +1,56 @@
+"""Test ast build pass module."""
+
+import io
+import os
+import sys
+
+from jaclang.compiler.program import JacProgram, py_code_gen
+from jaclang.utils.test import TestCase
+from jaclang.settings import settings
+from jaclang.compiler.passes.main import PreDynamoPass
+
+class PreDynamoPassTests(TestCase):
+    """Test pass module."""
+
+    TargetPass = PreDynamoPass
+
+    def setUp(self) -> None:
+        """Set up test."""
+        settings.predynamo_pass = True
+        return super().setUp()
+
+    def tearDown(self) -> None:
+        """Tear down test."""
+        settings.predynamo_pass = False
+        # Remove PreDynamoPass from global py_code_gen list if it was added
+        if PreDynamoPass in py_code_gen:
+            py_code_gen.remove(PreDynamoPass)
+        return super().tearDown()
+
+    def test_predynamo_where_assign(self) -> None:
+        """Test torch.where transformation."""
+        captured_output = io.StringIO()
+        sys.stdout = captured_output
+        settings.predynamo_pass = True
+        code_gen = JacProgram().compile(self.fixture_abs_path("predynamo_where_assign.jac"))
+        sys.stdout = sys.__stdout__
+        self.assertIn("torch.where", code_gen.unparse())
+        
+    def test_predynamo_where_return(self) -> None:
+        """Test torch.where transformation."""
+        captured_output = io.StringIO()
+        sys.stdout = captured_output
+        code_gen = JacProgram().compile(self.fixture_abs_path("predynamo_where_return.jac"))
+        sys.stdout = sys.__stdout__
+        self.assertIn("torch.where", code_gen.unparse())
+
+
+    def test_predynamo_fix3(self) -> None:
+        """Test torch.where transformation."""
+        captured_output = io.StringIO()
+        sys.stdout = captured_output
+        code_gen = JacProgram().compile(self.fixture_abs_path("predynamo_fix3.jac"))
+        sys.stdout = sys.__stdout__
+        unparsed_code = code_gen.unparse()
+        self.assertIn("__inv_freq = torch.where(", unparsed_code)
+        self.assertIn("self.register_buffer('inv_freq', __inv_freq, persistent=False);", unparsed_code)