Autograph works when applied to a function decorated with vmap

[rauletorresc]

Context: We want to ensure that autograph still works properly when applied to vmap.

Description of the Change: Handle callable objects at the transform method of the CatalystTransformer class.

[sc-60305]

[rauletorresc]

I think this code will never be reached, because qjit will realize earlier that the object is not a callable and will throw a different exception: https://github.com/PennyLaneAI/catalyst/blob/308f508e71a0af5e1aecc8560110deac2b5f84b1/frontend/catalyst/autograph/transformer.py#L61 That explains why I had to change the "unsupported object test" to "callable object test"

Do you have an idea which object type could hit that exception, so we keep testing unsupported objects? Or is it okay to skip it?

[codecov[bot]]

Codecov Report

Attention: Patch coverage is 75.00000% with 1 line in your changes missing coverage. Please review.

Project coverage is 97.92%. Comparing base (57b55c6) to head (7a15baa). Report is 1 commits behind head on main.

Files	Patch %	Lines
frontend/catalyst/autograph/transformer.py	75.00%	0 Missing and 1 partial :warning:

Additional details and impacted files

```diff @@ Coverage Diff @@ ## main #835 +/- ## ========================================== - Coverage 97.92% 97.92% -0.01% ========================================== Files 71 71 Lines 10240 10243 +3 Branches 1159 1161 +2 ========================================== + Hits 10028 10030 +2 Misses 169 169 - Partials 43 44 +1 ```

:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Have feedback on the report? Share it here.

[josh146]

Do you have an idea which object type could hit that exception, so we keep testing unsupported objects? Or is it okay to skip it?

@rauletorresc will defer to @dime10 here, but if that exception is no longer hit seems reasonable to not test it. Could conceivably turn it into an assertion rather than an else: raise ... statement.

[dime10]

Do you have an idea which object type could hit that exception, so we keep testing unsupported objects? Or is it okay to skip it?

@rauletorresc will defer to @dime10 here, but if that exception is no longer hit seems reasonable to not test it. Could conceivably turn it into an assertion rather than an else: raise ... statement.

Just any non-callable object I suppose, but it would have to pass through the previous init code. Looking at it, there is this function which is called first which already assumes a callable (has an assert inside): https://github.com/PennyLaneAI/catalyst/blob/e9c8eccdcedbeeb9b41acd17ddffd36148f61e5c/frontend/catalyst/jit.py#L489

Maybe we should raise an exception in QJIT.init if fn is not a callable? Then we can change it to assert inside the autograph code like Josh suggests.

Although having said that, we actually have this public function called run_autograph which runs AG without any QJIT involvement, so I think the exception should stay. You can use that function to test it though :)

[rauletorresc]

Although having said that, we actually have this public function called run_autograph which runs AG without any QJIT involvement, so I think the exception should stay. You can use that function to test it though :)

I went for this fix :)

[dime10]

Didn't realize this didn't make it into the release. Was there anything more that needed to be done?

[rauletorresc]

Didn't realize this didn't make it into the release. Was there anything more that needed to be done?

I am having a code coverage issue with a Lambda: https://app.codecov.io/gh/PennyLaneAI/catalyst/pull/835/blob/frontend/catalyst/autograph/transformer.py#L60

[dime10]

Didn't realize this didn't make it into the release. Was there anything more that needed to be done?

I am having a code coverage issue with a Lambda: https://app.codecov.io/gh/PennyLaneAI/catalyst/pull/835/blob/frontend/catalyst/autograph/transformer.py#L60

In general I don't think coverage should ever block an important contribution :)

For this PR, if we can't resolve why the line is partially uncovered I'm also happy to merge as is.

[dime10]

I just encountered another issue with autograph and vmap, and I was confused because I thought we had fixed it in this PR. This is the failing example:

@qml.qnode(qml.device("lightning.qubit", wires=20, shots=100))
def GHZState_if(use_clifford, clifford):
    if use_clifford:
        qml.QubitUnitary(clifford, wires=n_qubits - 1)
    return qml.probs()

qjit(vmap(GHZState_if, in_axes=(None, 0)), autograph=True)(True, jnp.zeros((2, 2, 2)))

Looking at the PR again it seems the test was inadequate because it doesn't involve any construct transformable by autograph that would indicate failure if the transform doesn't take place, or otherwise verifies the transform took place. I'll try to come up with another fix but it seems the lambda trick was not sufficient 🤔