PennyLaneAI / catalyst

A JIT compiler for hybrid quantum programs in PennyLane
https://docs.pennylane.ai/projects/catalyst
Apache License 2.0
122 stars 27 forks source link

[Bug] Compile time conditionals on Python types fail with Autograph #940

Closed josh146 closed 1 month ago

josh146 commented 1 month ago

The following function, which contains a conditional that is known at compile time, fails when Autograph is enabled:

@qml.qjit(autograph=True)
def f(x):
    n = 1

    if n:
        y = x ** 2
    else:
        y = x

    return y
>>> f(0.423)
CompileError: Compilation failed:
f:4:10: error: 'tensor.extract' op failed to verify that result type matches element type of tensor
    %1 = "tensor.extract"(%0) : (tensor<i64>) -> i1
         ^
f:4:10: note: see current operation: %1 = "tensor.extract"(%0) : (tensor<i64>) -> i1
Failed to parse module as MLIR source, retrying parsing as LLVM source
f: f:1:8: error: expected 'module asm'
module @f {
       ^
Failed to parse module as LLVM source

Note that properly converting the array to a boolean via jnp.all fixes this:

@qml.qjit(autograph=True)
def f(x):
    n = 1

    if jnp.all(n):
        y = x ** 2
    else:
        y = x

    return y
>>> f(0.423)
array(0.178929)

While this can be solved relatively easy if you recognize the MLIR error, this could be confusing to users.