fix fuse and args

Author: IvanGrigorik

pykokkos/core/fusion/fuse.py

@@ -1,5 +1,4 @@
 import ast
-import os
 from typing import Any, Dict, List, Set, Tuple, Union
 
 from .util import DeclarationsVisitor, VariableRenamer
@@ -47,7 +46,7 @@ def fuse_workunit_kwargs_and_params(
 
         for p in current_params:
             current_arg = current_kwargs[p.arg]
-            if "PK_FUSE_ARGS" in os.environ and id(current_arg) in view_ids:
+            if id(current_arg) in view_ids:
                 continue
 
             view_ids.add(id(current_arg))
@@ -131,7 +130,7 @@ def fuse_arguments(
                 continue
 
             current_arg = current_kwargs[old_name]
-            if "PK_FUSE_ARGS" in os.environ and id(current_arg) in fused_view_names:
+            if id(current_arg) in fused_view_names:
                 name_map[key] = fused_view_names[id(current_arg)]
                 continue
 

pykokkos/core/visitors/constructor_visitor.py

@@ -92,6 +92,15 @@ def visit_AnnAssign(
             if not isinstance(node.annotation, ast.Subscript):
                 return ()
 
+            if (
+                node.value is not None
+                and isinstance(node.value, ast.Subscript)
+                and isinstance(node.value.value, ast.Attribute)
+                and isinstance(node.value.value.value, ast.Name)
+                and node.value.value.value.id == "self"
+            ):
+                return (declref, None)
+
             decltype: cppast.ClassType = visitors_util.get_type(
                 node.annotation, self.pk_import
             )