Classical pre-processing not working when using grad with enzyme

error.txt

Traceback (most recent call last):
  File "/workspace/modules/catalyst/frontend/catalyst/compiler.py", line 365, in run_from_ir
    compiler_output = run_compiler_driver(
RuntimeError: Compilation failed:
f:6:3: error: operand #0 does not dominate this use
  func.func private @"<lambda>"(%arg0: tensor<f64>) -> tensor<f64> attributes {llvm.linkage = #llvm.linkage<internal>} {
  ^
f:6:3: note: diagnostic emitted with trace:
 #0 0x00007f08dcb4ac13 llvm::sys::PrintStackTrace(llvm::raw_ostream&, int) (/workspace/_build_dialects/python_packages/quantum/mlir_quantum/compiler_driver.so+0x8b4ac13)
 #1 0x00007f08da98c1b9 emitDiag(mlir::Location, mlir::DiagnosticSeverity, llvm::Twine const&) Diagnostics.cpp:0:0
 #2 0x00007f08da98c091 mlir::emitError(mlir::Location, llvm::Twine const&) (/workspace/_build_dialects/python_packages/quantum/mlir_quantum/compiler_driver.so+0x698c091)
 #3 0x00007f08da9f58fe mlir::Operation::emitError(llvm::Twine const&) (/workspace/_build_dialects/python_packages/quantum/mlir_quantum/compiler_driver.so+0x69f58fe)
 #4 0x00007f08daa17aa7 (anonymous namespace)::OperationVerifier::verifyDominanceOfContainedRegions(mlir::Operation&, mlir::DominanceInfo&) Verifier.cpp:0:0
 #5 0x00007f08daa17513 (anonymous namespace)::OperationVerifier::verifyOpAndDominance(mlir::Operation&) Verifier.cpp:0:0
 #6 0x00007f08daa16a90 (anonymous namespace)::OperationVerifier::verifyOpAndDominance(mlir::Operation&) Verifier.cpp:0:0
 #7 0x00007f08daa16133 mlir::verify(mlir::Operation*, bool) (/workspace/_build_dialects/python_packages/quantum/mlir_quantum/compiler_driver.so+0x6a16133)
 #8 0x00007f08da2b9d50 mlir::detail::OpToOpPassAdaptor::run(mlir::Pass*, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int) (/workspace/_build_dialects/python_packages/quantum/mlir_quantum/compiler_driver.so+0x62b9d50)
 #9 0x00007f08da2ba261 mlir::detail::OpToOpPassAdaptor::runPipeline(mlir::OpPassManager&, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int, mlir::PassInstrumentor*, mlir::PassInstrumentation::PipelineParentInfo const*) (/workspace/_build_dialects/python_packages/quantum/mlir_quantum/compiler_driver.so+0x62ba261)
#10 0x00007f08da2bc8c8 mlir::PassManager::run(mlir::Operation*) (/workspace/_build_dialects/python_packages/quantum/mlir_quantum/compiler_driver.so+0x62bc8c8)
#11 0x00007f08d7483214 runLowering(catalyst::driver::CompilerOptions const&, mlir::MLIRContext*, mlir::ModuleOp, catalyst::driver::CompilerOutput&) (/workspace/_build_dialects/python_packages/quantum/mlir_quantum/compiler_driver.so+0x3483214)
#12 0x00007f08d7483e38 QuantumDriverMain(catalyst::driver::CompilerOptions const&, catalyst::driver::CompilerOutput&) (/workspace/_build_dialects/python_packages/quantum/mlir_quantum/compiler_driver.so+0x3483e38)
#13 0x00007f08d747c7b6 void pybind11::cpp_function::initialize<pybind11_init_compiler_driver(pybind11::module_&)::$_8, std::unique_ptr<catalyst::driver::CompilerOutput, std::default_delete<catalyst::driver::CompilerOutput>>, char const*, char const*, char const*, bool, bool, pybind11::list, bool, pybind11::name, pybind11::scope, pybind11::sibling, pybind11::arg, pybind11::arg, pybind11::arg_v, pybind11::arg_v, pybind11::arg_v, pybind11::arg_v, pybind11::arg_v>(pybind11_init_compiler_driver(pybind11::module_&)::$_8&&, std::unique_ptr<catalyst::driver::CompilerOutput, std::default_delete<catalyst::driver::CompilerOutput>> (*)(char const*, char const*, char const*, bool, bool, pybind11::list, bool), pybind11::name const&, pybind11::scope const&, pybind11::sibling const&, pybind11::arg const&, pybind11::arg const&, pybind11::arg_v const&, pybind11::arg_v const&, pybind11::arg_v const&, pybind11::arg_v const&, pybind11::arg_v const&)::'lambda'(pybind11::detail::function_call&)::__invoke(pybind11::detail::function_call&) PyCompilerDriver.cpp:0:0
#14 0x00007f08d746de66 pybind11::cpp_function::dispatcher(_object*, _object*, _object*) PyCompilerDriver.cpp:0:0
#15 0x000055f2f02ad99e (/usr/bin/python3.10+0x15c99e)
#16 0x000055f2f02a44ab _PyObject_MakeTpCall (/usr/bin/python3.10+0x1534ab)
#17 0x000055f2f029d683 _PyEval_EvalFrameDefault (/usr/bin/python3.10+0x14c683)
#18 0x000055f2f02bbc31 (/usr/bin/python3.10+0x16ac31)
#19 0x000055f2f02bc8e2 PyObject_Call (/usr/bin/python3.10+0x16b8e2)
#20 0x000055f2f0298af0 _PyEval_EvalFrameDefault (/usr/bin/python3.10+0x147af0)
#21 0x000055f2f02ae1ec _PyFunction_Vectorcall (/usr/bin/python3.10+0x15d1ec)
#22 0x000055f2f02968cb _PyEval_EvalFrameDefault (/usr/bin/python3.10+0x1458cb)
#23 0x000055f2f02ae1ec _PyFunction_Vectorcall (/usr/bin/python3.10+0x15d1ec)
#24 0x000055f2f02968cb _PyEval_EvalFrameDefault (/usr/bin/python3.10+0x1458cb)
#25 0x000055f2f02bbd7e (/usr/bin/python3.10+0x16ad7e)
#26 0x000055f2f0298af0 _PyEval_EvalFrameDefault (/usr/bin/python3.10+0x147af0)
#27 0x000055f2f02a3634 _PyObject_FastCallDictTstate (/usr/bin/python3.10+0x152634)
#28 0x000055f2f02b8cac _PyObject_Call_Prepend (/usr/bin/python3.10+0x167cac)
#29 0x000055f2f03d6610 (/usr/bin/python3.10+0x285610)
#30 0x000055f2f02a44ab _PyObject_MakeTpCall (/usr/bin/python3.10+0x1534ab)
#31 0x000055f2f029c815 _PyEval_EvalFrameDefault (/usr/bin/python3.10+0x14b815)
#32 0x000055f2f0292ed6 (/usr/bin/python3.10+0x141ed6)
#33 0x000055f2f0389366 PyEval_EvalCode (/usr/bin/python3.10+0x238366)
#34 0x000055f2f03b6108 (/usr/bin/python3.10+0x265108)
#35 0x000055f2f03aef5b (/usr/bin/python3.10+0x25df5b)
#36 0x000055f2f03b5e55 (/usr/bin/python3.10+0x264e55)
#37 0x000055f2f03b5338 _PyRun_SimpleFileObject (/usr/bin/python3.10+0x264338)
#38 0x000055f2f03b5033 _PyRun_AnyFileObject (/usr/bin/python3.10+0x264033)
#39 0x000055f2f03a62de Py_RunMain (/usr/bin/python3.10+0x2552de)
#40 0x000055f2f037c32d Py_BytesMain (/usr/bin/python3.10+0x22b32d)
#41 0x00007f092b046d90 (/lib/x86_64-linux-gnu/libc.so.6+0x29d90)
#42 0x00007f092b046e40 __libc_start_main (/lib/x86_64-linux-gnu/libc.so.6+0x29e40)
#43 0x000055f2f037c225 _start (/usr/bin/python3.10+0x22b225)

f:6:3: note: see current operation: %0 = "func.call"(%3) <{callee = @g.pcount}> : (tensor<f64>) -> index
f:7:10: note: operand defined here (op in the same block)
    %0 = stablehlo.cosine %arg0 : tensor<f64>
         ^
While processing pipeline: QuantumCompilationPass
Failed to lower MLIR module


The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/workspace/src/compdriver/test3.py", line 22, in <module>
    print(f(0.14))
  File "/workspace/modules/catalyst/frontend/catalyst/compilation_pipelines.py", line 641, in __call__
    function, args = self._ensure_real_arguments_and_formal_parameters_are_compatible(
  File "/workspace/modules/catalyst/frontend/catalyst/compilation_pipelines.py", line 616, in _ensure_real_arguments_and_formal_parameters_are_compatible
    function = self.compile()
  File "/workspace/modules/catalyst/frontend/catalyst/compilation_pipelines.py", line 577, in compile
    shared_object, llvm_ir, inferred_func_data = self.compiler.run(self.mlir_module)
  File "/workspace/modules/catalyst/frontend/catalyst/compiler.py", line 410, in run
    return self.run_from_ir(
  File "/workspace/modules/catalyst/frontend/catalyst/compiler.py", line 375, in run_from_ir
    raise CompileError(*e.args) from e
catalyst.utils.exceptions.CompileError: Compilation failed:
f:6:3: error: operand #0 does not dominate this use
  func.func private @"<lambda>"(%arg0: tensor<f64>) -> tensor<f64> attributes {llvm.linkage = #llvm.linkage<internal>} {
  ^
f:6:3: note: diagnostic emitted with trace:
 #0 0x00007f08dcb4ac13 llvm::sys::PrintStackTrace(llvm::raw_ostream&, int) (/workspace/_build_dialects/python_packages/quantum/mlir_quantum/compiler_driver.so+0x8b4ac13)
 #1 0x00007f08da98c1b9 emitDiag(mlir::Location, mlir::DiagnosticSeverity, llvm::Twine const&) Diagnostics.cpp:0:0
 #2 0x00007f08da98c091 mlir::emitError(mlir::Location, llvm::Twine const&) (/workspace/_build_dialects/python_packages/quantum/mlir_quantum/compiler_driver.so+0x698c091)
 #3 0x00007f08da9f58fe mlir::Operation::emitError(llvm::Twine const&) (/workspace/_build_dialects/python_packages/quantum/mlir_quantum/compiler_driver.so+0x69f58fe)
 #4 0x00007f08daa17aa7 (anonymous namespace)::OperationVerifier::verifyDominanceOfContainedRegions(mlir::Operation&, mlir::DominanceInfo&) Verifier.cpp:0:0
 #5 0x00007f08daa17513 (anonymous namespace)::OperationVerifier::verifyOpAndDominance(mlir::Operation&) Verifier.cpp:0:0
 #6 0x00007f08daa16a90 (anonymous namespace)::OperationVerifier::verifyOpAndDominance(mlir::Operation&) Verifier.cpp:0:0
 #7 0x00007f08daa16133 mlir::verify(mlir::Operation*, bool) (/workspace/_build_dialects/python_packages/quantum/mlir_quantum/compiler_driver.so+0x6a16133)
 #8 0x00007f08da2b9d50 mlir::detail::OpToOpPassAdaptor::run(mlir::Pass*, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int) (/workspace/_build_dialects/python_packages/quantum/mlir_quantum/compiler_driver.so+0x62b9d50)
 #9 0x00007f08da2ba261 mlir::detail::OpToOpPassAdaptor::runPipeline(mlir::OpPassManager&, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int, mlir::PassInstrumentor*, mlir::PassInstrumentation::PipelineParentInfo const*) (/workspace/_build_dialects/python_packages/quantum/mlir_quantum/compiler_driver.so+0x62ba261)
#10 0x00007f08da2bc8c8 mlir::PassManager::run(mlir::Operation*) (/workspace/_build_dialects/python_packages/quantum/mlir_quantum/compiler_driver.so+0x62bc8c8)
#11 0x00007f08d7483214 runLowering(catalyst::driver::CompilerOptions const&, mlir::MLIRContext*, mlir::ModuleOp, catalyst::driver::CompilerOutput&) (/workspace/_build_dialects/python_packages/quantum/mlir_quantum/compiler_driver.so+0x3483214)
#12 0x00007f08d7483e38 QuantumDriverMain(catalyst::driver::CompilerOptions const&, catalyst::driver::CompilerOutput&) (/workspace/_build_dialects/python_packages/quantum/mlir_quantum/compiler_driver.so+0x3483e38)
#13 0x00007f08d747c7b6 void pybind11::cpp_function::initialize<pybind11_init_compiler_driver(pybind11::module_&)::$_8, std::unique_ptr<catalyst::driver::CompilerOutput, std::default_delete<catalyst::driver::CompilerOutput>>, char const*, char const*, char const*, bool, bool, pybind11::list, bool, pybind11::name, pybind11::scope, pybind11::sibling, pybind11::arg, pybind11::arg, pybind11::arg_v, pybind11::arg_v, pybind11::arg_v, pybind11::arg_v, pybind11::arg_v>(pybind11_init_compiler_driver(pybind11::module_&)::$_8&&, std::unique_ptr<catalyst::driver::CompilerOutput, std::default_delete<catalyst::driver::CompilerOutput>> (*)(char const*, char const*, char const*, bool, bool, pybind11::list, bool), pybind11::name const&, pybind11::scope const&, pybind11::sibling const&, pybind11::arg const&, pybind11::arg const&, pybind11::arg_v const&, pybind11::arg_v const&, pybind11::arg_v const&, pybind11::arg_v const&, pybind11::arg_v const&)::'lambda'(pybind11::detail::function_call&)::__invoke(pybind11::detail::function_call&) PyCompilerDriver.cpp:0:0
#14 0x00007f08d746de66 pybind11::cpp_function::dispatcher(_object*, _object*, _object*) PyCompilerDriver.cpp:0:0
#15 0x000055f2f02ad99e (/usr/bin/python3.10+0x15c99e)
#16 0x000055f2f02a44ab _PyObject_MakeTpCall (/usr/bin/python3.10+0x1534ab)
#17 0x000055f2f029d683 _PyEval_EvalFrameDefault (/usr/bin/python3.10+0x14c683)
#18 0x000055f2f02bbc31 (/usr/bin/python3.10+0x16ac31)
#19 0x000055f2f02bc8e2 PyObject_Call (/usr/bin/python3.10+0x16b8e2)
#20 0x000055f2f0298af0 _PyEval_EvalFrameDefault (/usr/bin/python3.10+0x147af0)
#21 0x000055f2f02ae1ec _PyFunction_Vectorcall (/usr/bin/python3.10+0x15d1ec)
#22 0x000055f2f02968cb _PyEval_EvalFrameDefault (/usr/bin/python3.10+0x1458cb)
#23 0x000055f2f02ae1ec _PyFunction_Vectorcall (/usr/bin/python3.10+0x15d1ec)
#24 0x000055f2f02968cb _PyEval_EvalFrameDefault (/usr/bin/python3.10+0x1458cb)
#25 0x000055f2f02bbd7e (/usr/bin/python3.10+0x16ad7e)
#26 0x000055f2f0298af0 _PyEval_EvalFrameDefault (/usr/bin/python3.10+0x147af0)
#27 0x000055f2f02a3634 _PyObject_FastCallDictTstate (/usr/bin/python3.10+0x152634)
#28 0x000055f2f02b8cac _PyObject_Call_Prepend (/usr/bin/python3.10+0x167cac)
#29 0x000055f2f03d6610 (/usr/bin/python3.10+0x285610)
#30 0x000055f2f02a44ab _PyObject_MakeTpCall (/usr/bin/python3.10+0x1534ab)
#31 0x000055f2f029c815 _PyEval_EvalFrameDefault (/usr/bin/python3.10+0x14b815)
#32 0x000055f2f0292ed6 (/usr/bin/python3.10+0x141ed6)
#33 0x000055f2f0389366 PyEval_EvalCode (/usr/bin/python3.10+0x238366)
#34 0x000055f2f03b6108 (/usr/bin/python3.10+0x265108)
#35 0x000055f2f03aef5b (/usr/bin/python3.10+0x25df5b)
#36 0x000055f2f03b5e55 (/usr/bin/python3.10+0x264e55)
#37 0x000055f2f03b5338 _PyRun_SimpleFileObject (/usr/bin/python3.10+0x264338)
#38 0x000055f2f03b5033 _PyRun_AnyFileObject (/usr/bin/python3.10+0x264033)
#39 0x000055f2f03a62de Py_RunMain (/usr/bin/python3.10+0x2552de)
#40 0x000055f2f037c32d Py_BytesMain (/usr/bin/python3.10+0x22b32d)
#41 0x00007f092b046d90 (/lib/x86_64-linux-gnu/libc.so.6+0x29d90)
#42 0x00007f092b046e40 __libc_start_main (/lib/x86_64-linux-gnu/libc.so.6+0x29e40)
#43 0x000055f2f037c225 _start (/usr/bin/python3.10+0x22b225)

f:6:3: note: see current operation: %0 = "func.call"(%3) <{callee = @g.pcount}> : (tensor<f64>) -> index
f:7:10: note: operand defined here (op in the same block)
    %0 = stablehlo.cosine %arg0 : tensor<f64>
         ^
While processing pipeline: QuantumCompilationPass
Failed to lower MLIR module