[py][nfc] imports clean up.

Cleans up imports throughout the python runtime. I trie following the principles in PEP8 (https://peps.python.org/pep-0008/#imports). The changes include: * Grouping imports two blocks: external and internal, w.r.t cudaq. * Not using relative paths when the module is in a parent dir (`..`) * Not using wildcard imports * Removing unused ones Signed-off-by: boschmitt <[email protected]>
NVIDIA · Feb 5, 2025 · 77ba55a · 77ba55a
1 parent a304c4f
commit 77ba55a
Show file tree

Hide file tree

Showing 25 changed files with 161 additions and 124 deletions.
diff --git a/python/cudaq/handlers/photonics_kernel.py b/python/cudaq/handlers/photonics_kernel.py
@@ -7,11 +7,10 @@
 # ============================================================================ #
 
 from __future__ import annotations
-
 from dataclasses import dataclass
 from typing import List
 
-from ..mlir._mlir_libs._quakeDialects import cudaq_runtime
+from cudaq.mlir._mlir_libs._quakeDialects import cudaq_runtime
 
 
 @dataclass

diff --git a/python/cudaq/kernel/analysis.py b/python/cudaq/kernel/analysis.py
@@ -6,11 +6,14 @@
 # the terms of the Apache License 2.0 which accompanies this distribution.     #
 # ============================================================================ #
 
-import ast, inspect, importlib, textwrap
+import ast
+import inspect
+import importlib
+import textwrap
+
+from cudaq.mlir._mlir_libs._quakeDialects import cudaq_runtime
+from cudaq.mlir.dialects import cc
 from .utils import globalAstRegistry, globalKernelRegistry, mlirTypeFromAnnotation
-from ..mlir.dialects import cc
-from ..mlir.ir import *
-from ..mlir._mlir_libs._quakeDialects import cudaq_runtime
 
 
 class MidCircuitMeasurementAnalyzer(ast.NodeVisitor):

diff --git a/python/cudaq/kernel/ast_bridge.py b/python/cudaq/kernel/ast_bridge.py
@@ -6,20 +6,36 @@
 # the terms of the Apache License 2.0 which accompanies this distribution.     #
 # ============================================================================ #
 import ast
-import hashlib
 import graphlib
-import sys, os
-from typing import Callable
-from collections import deque
 import numpy as np
+import os
+import sys
+from collections import deque
+
+from cudaq.mlir._mlir_libs._quakeDialects import (
+    cudaq_runtime, load_intrinsic, gen_vector_of_complex_constant,
+    register_all_dialects)
+from cudaq.mlir.dialects import arith, cc, complex, func, math, quake
+from cudaq.mlir.ir import (BoolAttr, Block, BlockArgument, Context, ComplexType,
+                           DenseI32ArrayAttr, DenseI64ArrayAttr, DictAttr,
+                           F32Type, F64Type, FlatSymbolRefAttr, FloatAttr,
+                           FunctionType, InsertionPoint, IntegerAttr,
+                           IntegerType, Location, Module, StringAttr,
+                           SymbolTable, TypeAttr, UnitAttr)
+from cudaq.mlir.passmanager import PassManager
 from .analysis import FindDepKernelsVisitor
-from .utils import globalAstRegistry, globalKernelRegistry, globalRegisteredOperations, nvqppPrefix, mlirTypeFromAnnotation, mlirTypeFromPyType, Color, mlirTypeToPyType, globalRegisteredTypes
-from ..mlir.ir import *
-from ..mlir.passmanager import *
-from ..mlir.dialects import quake, cc
-from ..mlir.dialects import builtin, func, arith, math, complex
-from ..mlir._mlir_libs._quakeDialects import cudaq_runtime, load_intrinsic, register_all_dialects, gen_vector_of_complex_constant
 from .captured_data import CapturedDataStorage
+from .utils import (
+    Color,
+    globalAstRegistry,
+    globalKernelRegistry,
+    globalRegisteredOperations,
+    globalRegisteredTypes,
+    nvqppPrefix,
+    mlirTypeFromAnnotation,
+    mlirTypeFromPyType,
+    mlirTypeToPyType,
+)
 
 State = cudaq_runtime.State
 

diff --git a/python/cudaq/kernel/captured_data.py b/python/cudaq/kernel/captured_data.py
@@ -9,14 +9,11 @@
 import uuid
 import numpy as np
 
-from ..mlir.ir import *
-from ..mlir.passmanager import *
-from ..mlir.execution_engine import *
-from ..mlir.dialects import quake, cc
-from ..mlir.dialects import builtin, func, arith
-from ..mlir._mlir_libs._quakeDialects import cudaq_runtime
-from ..mlir._mlir_libs._quakeDialects.cudaq_runtime import deletePointersToCudaqState
-from ..mlir._mlir_libs._quakeDialects.cudaq_runtime import deletePointersToStateData
+from cudaq.mlir._mlir_libs._quakeDialects import cudaq_runtime
+from cudaq.mlir.dialects import arith, cc, func
+from cudaq.mlir.ir import (ComplexType, F32Type, F64Type, FlatSymbolRefAttr,
+                           FunctionType, InsertionPoint, IntegerAttr,
+                           IntegerType, StringAttr, TypeAttr)
 
 kDynamicPtrIndex: int = -2147483648
 
@@ -48,8 +45,8 @@ def __del__(self):
         """
         Remove pointers to stored data for the current kernel.
         """
-        deletePointersToCudaqState(self.cudaqStateIDs)
-        deletePointersToStateData(self.arrayIDs)
+        cudaq_runtime.deletePointersToCudaqState(self.cudaqStateIDs)
+        cudaq_runtime.deletePointersToStateData(self.arrayIDs)
 
     def getIntegerAttr(self, type, value):
         """

diff --git a/python/cudaq/kernel/kernel_builder.py b/python/cudaq/kernel/kernel_builder.py
@@ -6,28 +6,49 @@
 # the terms of the Apache License 2.0 which accompanies this distribution.     #
 # ============================================================================ #
 
-from functools import partialmethod
-import hashlib
-import uuid
+import numpy as np
 import random
 import re
 import string
-import sys
-import numpy as np
+from functools import partialmethod
 from typing import get_origin
-from .quake_value import QuakeValue
-from .kernel_decorator import PyKernelDecorator
-from .utils import mlirTypeFromPyType, nvqppPrefix, emitFatalError, emitWarning, mlirTypeToPyType, emitErrorIfInvalidPauli, globalRegisteredOperations
-from .common.givens import givens_builder
-from .common.fermionic_swap import fermionic_swap_builder
-from .captured_data import CapturedDataStorage
 
-from ..mlir.ir import *
-from ..mlir.passmanager import *
-from ..mlir.execution_engine import *
-from ..mlir.dialects import quake, cc
-from ..mlir.dialects import builtin, func, arith, math, complex as complexDialect
-from ..mlir._mlir_libs._quakeDialects import cudaq_runtime, load_intrinsic, register_all_dialects, gen_vector_of_complex_constant
+from cudaq.mlir.ir import (
+    BoolAttr,
+    Block,
+    ComplexType,
+    Context,
+    DenseI32ArrayAttr,
+    DictAttr,
+    F32Type,
+    F64Type,
+    FlatSymbolRefAttr,
+    FloatAttr,
+    InsertionPoint,
+    IntegerAttr,
+    IntegerType,
+    Location,
+    Module,
+    StringAttr,
+    SymbolTable,
+    TypeAttr,
+    UnitAttr,
+)
+from cudaq.mlir.passmanager import PassManager
+from cudaq.mlir.execution_engine import ExecutionEngine
+from cudaq.mlir.dialects import (complex as complexDialect, arith, quake, cc,
+                                 func, math)
+from cudaq.mlir._mlir_libs._quakeDialects import (
+    cudaq_runtime, gen_vector_of_complex_constant, load_intrinsic,
+    register_all_dialects)
+from .captured_data import CapturedDataStorage
+from .common.fermionic_swap import fermionic_swap_builder
+from .common.givens import givens_builder
+from .kernel_decorator import PyKernelDecorator
+from .quake_value import QuakeValue
+from .utils import (emitErrorIfInvalidPauli, emitFatalError, emitWarning,
+                    globalRegisteredOperations, mlirTypeFromPyType,
+                    mlirTypeToPyType, nvqppPrefix)
 
 kDynamicPtrIndex: int = -2147483648
 

diff --git a/python/cudaq/kernel/kernel_decorator.py b/python/cudaq/kernel/kernel_decorator.py
@@ -5,22 +5,23 @@
 # This source code and the accompanying materials are made available under     #
 # the terms of the Apache License 2.0 which accompanies this distribution.     #
 # ============================================================================ #
-import ast, sys, traceback
+import ast
 import importlib
 import inspect
 import json
-from typing import Callable
-from ..mlir.ir import *
-from ..mlir.passmanager import *
-from ..mlir.dialects import quake, cc, func
-from .ast_bridge import compile_to_mlir, PyASTBridge
-from .utils import mlirTypeFromPyType, nvqppPrefix, mlirTypeToPyType, globalAstRegistry, emitFatalError, emitErrorIfInvalidPauli, globalRegisteredTypes
+import numpy as np
+
+from cudaq.handlers import PhotonicsHandler
+from cudaq.mlir._mlir_libs._quakeDialects import cudaq_runtime
+from cudaq.mlir.dialects import cc, func
+from cudaq.mlir.ir import (ComplexType, F32Type, F64Type, IntegerType,
+                           SymbolTable)
 from .analysis import MidCircuitMeasurementAnalyzer, HasReturnNodeVisitor
-from ..mlir._mlir_libs._quakeDialects import cudaq_runtime
+from .ast_bridge import compile_to_mlir, PyASTBridge
 from .captured_data import CapturedDataStorage
-from ..handlers import PhotonicsHandler
-
-import numpy as np
+from .utils import (emitFatalError, emitErrorIfInvalidPauli, globalAstRegistry,
+                    globalRegisteredTypes, mlirTypeFromPyType, mlirTypeToPyType,
+                    nvqppPrefix)
 
 # This file implements the decorator mechanism needed to
 # JIT compile CUDA-Q kernels. It exposes the cudaq.kernel()

diff --git a/python/cudaq/kernel/quake_value.py b/python/cudaq/kernel/quake_value.py
@@ -5,19 +5,11 @@
 # This source code and the accompanying materials are made available under     #
 # the terms of the Apache License 2.0 which accompanies this distribution.     #
 # ============================================================================ #
-from functools import partialmethod
-import inspect
-import sys
-import random
-import string
-import numpy as np
-import ctypes
-
-from ..mlir.ir import *
-from ..mlir.passmanager import *
-from ..mlir.dialects import quake, cc
-from ..mlir.dialects import builtin, func, arith
-from ..mlir._mlir_libs._quakeDialects import cudaq_runtime
+
+from cudaq.mlir._mlir_libs._quakeDialects import cudaq_runtime
+from cudaq.mlir.dialects import arith, quake, cc
+from cudaq.mlir.ir import (DenseI32ArrayAttr, F64Type, FloatAttr, Location,
+                           IntegerAttr, IntegerType)
 from .utils import mlirTypeFromPyType
 
 qvector = cudaq_runtime.qvector

diff --git a/python/cudaq/kernel/register_op.py b/python/cudaq/kernel/register_op.py
@@ -6,14 +6,14 @@
 # the terms of the Apache License 2.0 which accompanies this distribution.     #
 # ============================================================================ #
 
-from functools import partialmethod
 import math
 import numpy as np
+from functools import partialmethod
 from typing import Callable, List
 
-from .utils import globalRegisteredOperations
+from cudaq.mlir._mlir_libs._quakeDialects import cudaq_runtime
 from .kernel_builder import PyKernel, __generalCustomOperation
-from ..mlir._mlir_libs._quakeDialects import cudaq_runtime
+from .utils import globalRegisteredOperations
 
 
 def register_operation(operation_name: str, unitary):

diff --git a/python/cudaq/kernel/utils.py b/python/cudaq/kernel/utils.py
@@ -6,16 +6,16 @@
 # the terms of the Apache License 2.0 which accompanies this distribution.     #
 # ============================================================================ #
 from __future__ import annotations
-
-from ..mlir._mlir_libs._quakeDialects import cudaq_runtime
-from ..mlir.dialects import quake, cc
-from ..mlir.ir import *
-from ..mlir.passmanager import *
-import numpy as np
-from typing import Callable, List
-import ast, sys, traceback
+import ast
 import re
-from typing import get_origin
+import sys
+import traceback
+import numpy as np
+from typing import get_origin, Callable, List
+
+from cudaq.mlir._mlir_libs._quakeDialects import cudaq_runtime
+from cudaq.mlir.dialects import quake, cc
+from cudaq.mlir.ir import ComplexType, F32Type, F64Type, IntegerType
 
 State = cudaq_runtime.State
 qvector = cudaq_runtime.qvector

diff --git a/python/cudaq/kernels/hwe.py b/python/cudaq/kernels/hwe.py
@@ -5,8 +5,6 @@
 # This source code and the accompanying materials are made available under     #
 # the terms of the Apache License 2.0 which accompanies this distribution.     #
 # ============================================================================ #
-import numpy as np
-import cudaq
 
 
 def num_hwe_parameters(numQubits, numLayers):

diff --git a/python/cudaq/operator/cudm_helpers.py b/python/cudaq/operator/cudm_helpers.py
@@ -6,12 +6,13 @@
 # the terms of the Apache License 2.0 which accompanies this distribution.     #
 # ============================================================================ #
 
+import logging
 import numpy
-from typing import Any, Mapping, List, Sequence, Union
 from numbers import Number
+from typing import Any, Mapping, List, Sequence, Union
+
 from .expressions import ElementaryOperator, ScalarOperator
 from .manipulation import OperatorArithmetics
-import logging
 from .schedule import Schedule
 
 cudm = None

diff --git a/python/cudaq/operator/cudm_solver.py b/python/cudaq/operator/cudm_solver.py
@@ -10,10 +10,9 @@
 from typing import Sequence, Mapping, List, Optional
 
 from .cudm_helpers import CuDensityMatOpConversion, constructLiouvillian
-from ..runtime.observe import observe
 from .schedule import Schedule
 from .expressions import Operator
-from ..mlir._mlir_libs._quakeDialects import cudaq_runtime
+from cudaq.mlir._mlir_libs._quakeDialects import cudaq_runtime
 from .cudm_helpers import cudm, CudmStateType
 from .cudm_state import CuDensityMatState, as_cudm_state
 from .helpers import InitialState, InitialStateArgT

diff --git a/python/cudaq/operator/cudm_state.py b/python/cudaq/operator/cudm_state.py
@@ -6,12 +6,14 @@
 # the terms of the Apache License 2.0 which accompanies this distribution.     #
 # ============================================================================ #
 
-from cuquantum.densitymat import DenseMixedState, DensePureState, WorkStream
-import numpy, cupy, atexit
-from typing import Sequence
+import cupy
+import numpy
 from cupy.cuda.memory import MemoryPointer, UnownedMemory
-from ..mlir._mlir_libs._quakeDialects import cudaq_runtime
 from cuquantum.bindings import cudensitymat as cudm
+from cuquantum.densitymat import DenseMixedState, DensePureState, WorkStream
+from typing import Sequence
+
+from cudaq.mlir._mlir_libs._quakeDialects import cudaq_runtime
 from .helpers import InitialState
 
 

diff --git a/python/cudaq/operator/definitions.py b/python/cudaq/operator/definitions.py
@@ -6,13 +6,15 @@
 # the terms of the Apache License 2.0 which accompanies this distribution.     #
 # ============================================================================ #
 
-import numpy, scipy  # type: ignore
+import numpy
+import scipy
 from numpy.typing import NDArray
 from typing import Sequence
 
+from cudaq.mlir._mlir_libs._quakeDialects import cudaq_runtime
 from .helpers import NumericType
-from .expressions import OperatorSum, ProductOperator, ElementaryOperator, ScalarOperator, RydbergHamiltonian
-from ..mlir._mlir_libs._quakeDialects import cudaq_runtime
+from .expressions import (OperatorSum, ProductOperator, ElementaryOperator,
+                          ScalarOperator)
 
 
 # Operators as defined here (watch out of differences in convention):

diff --git a/python/cudaq/operator/evolution.py b/python/cudaq/operator/evolution.py
@@ -7,23 +7,22 @@
 # ============================================================================ #
 
 from __future__ import annotations
-import numpy, scipy, sys, uuid
-from numpy.typing import NDArray
-from typing import Callable, Iterable, Mapping, Optional, Sequence
-import json
+import numpy
 import random
 import string
 import warnings
+import uuid
+from numpy.typing import NDArray
+from typing import Callable, Iterable, Mapping, Optional, Sequence
 
+from cudaq.kernel.kernel_builder import PyKernel, make_kernel
+from cudaq.kernel.register_op import register_operation
+from cudaq.kernel.utils import ahkPrefix
+from cudaq.mlir._mlir_libs._quakeDialects import cudaq_runtime
 from .expressions import Operator, RydbergHamiltonian
 from .helpers import NumericType, InitialState, InitialStateArgT
 from .integrator import BaseIntegrator
 from .schedule import Schedule
-from ..kernel.register_op import register_operation
-from ..kernel.utils import ahkPrefix
-from ..mlir._mlir_libs._quakeDialects import cudaq_runtime
-from ..kernel.kernel_builder import PyKernel, make_kernel
-from ..runtime.observe import observe
 
 
 def _taylor_series_expm(op_matrix: NDArray[numpy.complexfloating],