Hello there,
I have an Issue with the jit computation if gradients in a (in principle) hybrid model using Jax and Catalyst. The following ends up with an error regarding the compiler (I guess).
Is there some advice what I can do about this?
Best regards
minn-bj
Code:
import pennylane as qml
from jax import numpy as np
import catalyst
# Hyperparameters
compiler="catalyst"
device="lightning.kokkos"
n_wires = 22
# Sample
data = qml.numpy.random.rand(n_wires)
# Weights
weights = np.ones([n_wires])
# Device
dev = qml.device(device, wires=n_wires)
# Quantum circuit
@qml.qnode(dev, interface="jax", diff_method="adjoint")
def circuit(data, weights):
for i in range(n_wires):
qml.RY(data[i], wires=i)
qml.RX(weights[i], wires=i)
return [qml.expval(qml.PauliZ(i)) for i in range(n_wires)]
# JIT circuit
circuit = qml.qjit(circuit, compiler=compiler)
# Loss
def loss_fn(weights, data):
predictions = np.array(circuit(data, weights))
loss = np.average((data - predictions) ** 2)
return loss
# Loss Gradient
grad = qml.qjit(catalyst.grad(loss_fn, method="auto"))
# Compute Grad
print(grad(weights, data))
Some more information about the system (sorry for the format):
Name: pennylane Version: 0.43.0 Summary: PennyLane is a cross-platform Python library for quantum computing, quantum machine learning, and quantum chemistry. Train a quantum computer the same way as a neural network. Home-page: Author: Author-email: License: Location: [/opt/conda/lib/python3.11/site-packages](https://vscode-remote+qcokain-002eintra-002edlr-002ede.vscode-resource.vscode-cdn.net/opt/conda/lib/python3.11/site-packages) Requires: appdirs, autograd, autoray, cachetools, diastatic-malt, networkx, numpy, packaging, pennylane-lightning, requests, rustworkx, scipy, tomlkit, typing_extensions Required-by: pennylane_catalyst, pennylane_lightning, pennylane_lightning_gpu, pennylane_lightning_kokkos Platform info: Linux-5.15.0-151-generic-x86_64-with-glibc2.35 Python version: 3.11.10 Numpy version: 2.4.0 Scipy version: 1.16.3 JAX version: 0.6.2 Installed devices: - nvidia.custatevec (pennylane_catalyst-0.13.0) - nvidia.cutensornet (pennylane_catalyst-0.13.0) - oqc.cloud (pennylane_catalyst-0.13.0) - softwareq.qpp (pennylane_catalyst-0.13.0) - lightning.qubit (pennylane_lightning-0.43.0) - default.clifford (pennylane-0.43.0) - default.gaussian (pennylane-0.43.0) - default.mixed (pennylane-0.43.0) - default.qubit (pennylane-0.43.0) - default.qutrit (pennylane-0.43.0) - default.qutrit.mixed (pennylane-0.43.0) - default.tensor (pennylane-0.43.0) - null.qubit (pennylane-0.43.0) - reference.qubit (pennylane-0.43.0) - lightning.gpu (pennylane_lightning_gpu-0.43.0) - lightning.kokkos (pennylane_lightning_kokkos-0.43.0)
Error msg:
CalledProcessError Traceback (most recent call last)
File /opt/conda/lib/python3.11/site-packages/catalyst/compiler.py:448, in Compiler.run_from_ir(self, ir, module_name, workspace)
447 print(f"\[SYSTEM\] {’ '.join(cmd)}", file=self.options.logfile)
→ 448 result = subprocess.run(cmd, check=True, capture_output=True, text=True)
449 if self.options.verbose or os.getenv(“ENABLE_DIAGNOSTICS”):
File /opt/conda/lib/python3.11/subprocess.py:571, in run(input, capture_output, timeout, check, \*popenargs, \*\*kwargs)
570 if check and retcode:
→ 571 raise CalledProcessError(retcode, process.args,
572 output=stdout, stderr=stderr)
573 return CompletedProcess(process.args, retcode, stdout, stderr)
CalledProcessError: Command ‘\[’/opt/conda/bin/catalyst’, ‘-o’, ‘/tmp/grad.loss_fnbi25h3bs/grad.loss_fn.ll’, ‘–module-name’, ‘grad.loss_fn’, ‘–workspace’, ‘/tmp/grad.loss_fnbi25h3bs’, ‘-verify-each=false’, ‘–catalyst-pipeline’, ‘EnforceRuntimeInvariantsPass(split-multiple-tapes;builtin.module(apply-transform-sequence);inline-nested-module),HLOLoweringPass(canonicalize;func.func(chlo-legalize-to-stablehlo);func.func(stablehlo-legalize-control-flow);func.func(stablehlo-aggressive-simplification);stablehlo-legalize-to-linalg;func.func(stablehlo-legalize-to-std);func.func(stablehlo-legalize-sort);stablehlo-convert-to-signless;canonicalize;scatter-lowering;hlo-custom-call-lowering;cse;func.func(linalg-detensorize{aggressive-mode});detensorize-scf;detensorize-function-boundary;canonicalize;symbol-dce),QuantumCompilationPass(annotate-function;lower-mitigation;lower-gradients;adjoint-lowering),BufferizationPass(inline;convert-tensor-to-linalg;convert-elementwise-to-linalg;gradient-preprocess;one-shot-bufferize{bufferize-function-boundaries allow-return-allocs-from-loops function-boundary-type-conversion=identity-layout-map unknown-type-conversion=identity-layout-map};canonicalize;gradient-postprocess;func.func(buffer-hoisting);func.func(buffer-loop-hoisting);func.func(promote-buffers-to-stack);func.func(buffer-deallocation);convert-arraylist-to-memref;convert-bufferization-to-memref;canonicalize;cp-global-memref),MLIRToLLVMDialect(expand-realloc;convert-gradient-to-llvm;memrefcpy-to-linalgcpy;func.func(convert-linalg-to-loops);convert-scf-to-cf;expand-strided-metadata;lower-affine;arith-expand;convert-complex-to-standard;convert-complex-to-llvm;convert-math-to-llvm;convert-math-to-libm;convert-arith-to-llvm;memref-to-llvm-tbaa;finalize-memref-to-llvm{use-generic-functions};convert-index-to-llvm;convert-catalyst-to-llvm;convert-quantum-to-llvm;emit-catalyst-py-interface;canonicalize;reconcile-unrealized-casts;gep-inbounds;register-inactive-callback),’, ‘/tmp/grad.loss_fnbi25h3bs/tmp76lhqb8r.mlir’\]’ died with <Signals.SIGABRT: 6>.
The above exception was the direct cause of the following exception:
CompileError Traceback (most recent call last)
Cell In\[5\], line 42
39 grad = qml.qjit(catalyst.grad(loss_fn, method=“auto”))
41 # Compute Grad
—> 42 print(grad(weights, data))
File /opt/conda/lib/python3.11/site-packages/pennylane/logging/decorators.py:61, in log_string_debug_func..wrapper_entry(\*args, \*\*kwargs)
54 s_caller = “::L”.join(
55 \[str(i) for i in inspect.getouterframes(inspect.currentframe(), 2)\[1\]\[1:3\]\]
56 )
57 lgr.debug(
58 f"Calling {f_string} from {s_caller}",
59 \*\*\_debug_log_kwargs,
60 )
—> 61 return func(\*args, \*\*kwargs)
File /opt/conda/lib/python3.11/site-packages/catalyst/jit.py:590, in QJIT.**call**(self, \*args, \*\*kwargs)
586 kwargs = {“static_argnums”: self.compile_options.static_argnums, \*\*kwargs}
588 return self.user_function(\*args, \*\*kwargs)
→ 590 requires_promotion = self.jit_compile(args, \*\*kwargs)
592 # If we receive tracers as input, dispatch to the JAX integration.
593 if any(isinstance(arg, jax.core.Tracer) for arg in tree_flatten(args)\[0\]):
File /opt/conda/lib/python3.11/site-packages/pennylane/logging/decorators.py:61, in log_string_debug_func..wrapper_entry(\*args, \*\*kwargs)
54 s_caller = “::L”.join(
55 \[str(i) for i in inspect.getouterframes(inspect.currentframe(), 2)\[1\]\[1:3\]\]
56 )
57 lgr.debug(
58 f"Calling {f_string} from {s_caller}",
59 \*\*\_debug_log_kwargs,
60 )
—> 61 return func(\*args, \*\*kwargs)
File /opt/conda/lib/python3.11/site-packages/catalyst/jit.py:666, in QJIT.jit_compile(self, args, \*\*kwargs)
663 self.jaxpr, self.out_type, self.out_treedef, self.c_sig = self.capture(args, \*\*kwargs)
665 self.mlir_module = self.generate_ir()
→ 666 self.compiled_function, \_ = self.compile()
668 self.fn_cache.insert(self.compiled_function, args, self.out_treedef, self.workspace)
670 elif self.compiled_function is not cached_fn.compiled_fn:
671 # Restore active state from cache.
File /opt/conda/lib/python3.11/site-packages/catalyst/debug/instruments.py:145, in instrument..wrapper(\*args, \*\*kwargs)
142 @functools.wraps(fn)
143 def wrapper(\*args, \*\*kwargs):
144 if not InstrumentSession.active:
→ 145 return fn(\*args, \*\*kwargs)
147 with ResultReporter(stage_name, has_finegrained) as reporter:
148 self = args\[0\]
File /opt/conda/lib/python3.11/site-packages/pennylane/logging/decorators.py:61, in log_string_debug_func..wrapper_entry(\*args, \*\*kwargs)
54 s_caller = “::L”.join(
55 \[str(i) for i in inspect.getouterframes(inspect.currentframe(), 2)\[1\]\[1:3\]\]
56 )
57 lgr.debug(
58 f"Calling {f_string} from {s_caller}",
59 \*\*\_debug_log_kwargs,
60 )
—> 61 return func(\*args, \*\*kwargs)
File /opt/conda/lib/python3.11/site-packages/catalyst/jit.py:816, in QJIT.compile(self)
810 shared_object, llvm_ir = self.compiler.run_from_ir(
811 self.overwrite_ir,
812 str(self.mlir_module.operation.attributes\[“sym_name”\]).replace(‘"’, “”),
813 self.workspace,
814 )
815 else:
→ 816 shared_object, llvm_ir = self.compiler.run(self.mlir_module, self.workspace)
818 compiled_fn = CompiledFunction(
819 shared_object, func_name, restype, self.out_type, self.compile_options
820 )
822 return compiled_fn, llvm_ir
File /opt/conda/lib/python3.11/site-packages/pennylane/logging/decorators.py:61, in log_string_debug_func..wrapper_entry(\*args, \*\*kwargs)
54 s_caller = “::L”.join(
55 \[str(i) for i in inspect.getouterframes(inspect.currentframe(), 2)\[1\]\[1:3\]\]
56 )
57 lgr.debug(
58 f"Calling {f_string} from {s_caller}",
59 \*\*\_debug_log_kwargs,
60 )
—> 61 return func(\*args, \*\*kwargs)
File /opt/conda/lib/python3.11/site-packages/catalyst/compiler.py:522, in Compiler.run(self, mlir_module, \*args, \*\*kwargs)
519 compiler = PythonCompiler()
520 mlir_module = compiler.run(mlir_module)
→ 522 return self.run_from_ir(
523 mlir_module.operation.get_asm(
524 binary=False, print_generic_op_form=False, assume_verified=True
525 ),
526 str(mlir_module.operation.attributes\[“sym_name”\]).replace(‘"’, “”),
527 \*args,
528 \*\*kwargs,
529 )
File /opt/conda/lib/python3.11/site-packages/pennylane/logging/decorators.py:61, in log_string_debug_func..wrapper_entry(\*args, \*\*kwargs)
54 s_caller = “::L”.join(
55 \[str(i) for i in inspect.getouterframes(inspect.currentframe(), 2)\[1\]\[1:3\]\]
56 )
57 lgr.debug(
58 f"Calling {f_string} from {s_caller}",
59 \*\*\_debug_log_kwargs,
60 )
—> 61 return func(\*args, \*\*kwargs)
File /opt/conda/lib/python3.11/site-packages/catalyst/compiler.py:455, in Compiler.run_from_ir(self, ir, module_name, workspace)
453 print(result.stderr.strip(), file=self.options.logfile)
454 except subprocess.CalledProcessError as e: # pragma: nocover
→ 455 raise CompileError(f"catalyst failed with error code {e.returncode}: {e.stderr}") from e
457 if os.path.exists(output_ir_name):
458 with open(output_ir_name, “r”, encoding=“utf-8”) as f:
CompileError: catalyst failed with error code -6: catalyst: /\__w/catalyst/catalyst/mlir/lib/Gradient/Transforms/ConversionPatterns.cpp:225: virtual llvm::LogicalResult {anonymous}::AdjointOpPattern::matchAndRewrite(catalyst::gradient::AdjointOp, catalyst::gradient::AdjointOpAdaptor, mlir::ConversionPatternRewriter&) const: Assertion `callee && callee.getNumResults() == 2 && "invalid qfunc symbol in adjoint op"' failed. PLEASE submit a bug report to https://github.com/llvm/llvm-project/issues/ and include the crash backtrace. Stack dump: 0. Program arguments: /opt/conda/bin/catalyst -o /tmp/grad.loss_fnbi25h3bs/grad.loss_fn.ll --module-name grad.loss_fn --workspace /tmp/grad.loss_fnbi25h3bs -verify-each=false --catalyst-pipeline "EnforceRuntimeInvariantsPass(split-multiple-tapes;builtin.module(apply-transform-sequence);inline-nested-module),HLOLoweringPass(canonicalize;func.func(chlo-legalize-to-stablehlo);func.func(stablehlo-legalize-control-flow);func.func(stablehlo-aggressive-simplification);stablehlo-legalize-to-linalg;func.func(stablehlo-legalize-to-std);func.func(stablehlo-legalize-sort);stablehlo-convert-to-signless;canonicalize;scatter-lowering;hlo-custom-call-lowering;cse;func.func(linalg-detensorize{aggressive-mode});detensorize-scf;detensorize-function-boundary;canonicalize;symbol-dce),QuantumCompilationPass(annotate-function;lower-mitigation;lower-gradients;adjoint-lowering),BufferizationPass(inline;convert-tensor-to-linalg;convert-elementwise-to-linalg;gradient-preprocess;one-shot-bufferize{bufferize-function-boundaries allow-return-allocs-from-loops function-boundary-type-conversion=identity-layout-map unknown-type-conversion=identity-layout-map};canonicalize;gradient-postprocess;func.func(buffer-hoisting);func.func(buffer-loop-hoisting);func.func(promote-buffers-to-stack);func.func(buffer-deallocation);convert-arraylist-to-memref;convert-bufferization-to-memref;canonicalize;cp-global-memref),MLIRToLLVMDialect(expand-realloc;convert-gradient-to-llvm;memrefcpy-to-linalgcpy;func.func(convert-linalg-to-loops);convert-scf-to-cf;expand-strided-metadata;lower-affine;arith-expand;convert-complex-to-standard;convert-complex-to-llvm;convert-math-to-llvm;convert-math-to-libm;convert-arith-to-llvm;memref-to-llvm-tbaa;finalize-memref-to-llvm{use-generic-functions};convert-index-to-llvm;convert-catalyst-to-llvm;convert-quantum-to-llvm;emit-catalyst-py-interface;canonicalize;reconcile-unrealized-casts;gep-inbounds;register-inactive-callback)," /tmp/grad.loss_fnbi25h3bs/tmp76lhqb8r.mlir Stack dump without symbol names (ensure you have llvm-symbolizer in your PATH or set the environment var `LLVM_SYMBOLIZER_PATH\` to point to it):
0 catalyst 0x000000000afba3eb llvm::sys::PrintStackTrace(llvm::raw_ostream&, int) + 43
1 catalyst 0x000000000afb781b llvm::sys::RunSignalHandlers() + 43
2 catalyst 0x000000000afb7941
3 libc.so.6 0x00007fcd5f66c520
4 libc.so.6 0x00007fcd5f6c09fc pthread_kill + 300
5 libc.so.6 0x00007fcd5f66c476 raise + 22
6 libc.so.6 0x00007fcd5f6527f3 abort + 211
7 libc.so.6 0x00007fcd5f65271b
8 libc.so.6 0x00007fcd5f663e96
9 catalyst 0x00000000070be160
10 catalyst 0x00000000070b1f23
11 catalyst 0x000000000ab4f19b mlir::ConversionPattern::matchAndRewrite(mlir::Operation\*, mlir::PatternRewriter&) const + 539
12 catalyst 0x000000000abab5ea mlir::PatternApplicator::matchAndRewrite(mlir::Operation\*, mlir::PatternRewriter&, llvm::function_ref<bool (mlir::Pattern const&)>, llvm::function_ref<void (mlir::Pattern const&)>, llvm::function_ref<llvm::LogicalResult (mlir::Pattern const&)>) + 1674
13 catalyst 0x000000000ab515bc
14 catalyst 0x000000000ab51a57 mlir::OperationConverter::convert(mlir::ConversionPatternRewriter&, mlir::Operation\*) + 39
15 catalyst 0x000000000ab5630c mlir::OperationConverter::convertOperations(llvm::ArrayRef<mlir::Operation*>) + 284
16 catalyst 0x000000000ab58862 mlir::applyPartialConversion(mlir::Operation\*, mlir::ConversionTarget const&, mlir::FrozenRewritePatternSet const&, mlir::ConversionConfig) + 130
17 catalyst 0x000000000713d0eb
18 catalyst 0x000000000abe4b56 mlir::detail::OpToOpPassAdaptor::run(mlir::Pass\*, mlir::Operation\*, mlir::AnalysisManager, bool, unsigned int) + 1126
19 catalyst 0x000000000abe4f60 mlir::detail::OpToOpPassAdaptor::runPipeline(mlir::OpPassManager&, mlir::Operation\*, mlir::AnalysisManager, bool, unsigned int, mlir::PassInstrumentor\*, mlir::PassInstrumentation::PipelineParentInfo const\*) + 336
20 catalyst 0x000000000abe5e69 mlir::PassManager::run(mlir::Operation\*) + 1193
21 catalyst 0x0000000003fbdffb runPipeline(mlir::PassManager&, catalyst::driver::CompilerOptions const&, catalyst::driver::CompilerOutput&, catalyst::driver::Pipeline&, bool, mlir::ModuleOp) + 251
22 catalyst 0x0000000003fbe84a runLowering(catalyst::driver::CompilerOptions const&, mlir::MLIRContext\*, mlir::ModuleOp, catalyst::driver::CompilerOutput&, mlir::TimingScope&) + 1530
23 catalyst 0x0000000003fc0c55 QuantumDriverMain(catalyst::driver::CompilerOptions const&, catalyst::driver::CompilerOutput&, mlir::DialectRegistry&) + 6645
24 catalyst 0x0000000003fc4a1f QuantumDriverMainFromCL(int, char\*\*) + 9391
25 libc.so.6 0x00007fcd5f653d90
26 libc.so.6 0x00007fcd5f653e40 \__libc_start_main + 128
27 catalyst 0x0000000003fa01ee \_start + 46
