Overusing Catalyst callbacks

Catalyst
1

Hihi! I’m exploring an application for Catalyst callbacks where I need to call a Python function that doesn’t particularly fit into the ‘pure_callback’ requirement from the docs. Concretely, my function accesses a global object living in Python-controlled space and returns different results from the same inputs because it depends on the global state.

Surprisingly, it works quite well.

My question is: in what scenarios should I expect impure callbacks to cause failures? Can I still hope for correct execution if I don’t use differentiation? Or if I only run circuit in one-shot mcm mode?

Thanks!
PS

Consider the following code as an illustration of callback overuse

def test_pl_random_callback_works():
    prob=0.5
    num_shots=1000

    @pure_callback
    def random_callback(key, prob:ShapedArray([],dtype=float)) -> ShapedArray([],dtype=bool):
        res_bool = random.random() < float(prob)
        return jnp.array(res_bool, dtype=bool)

    @qjit
    @qml.qnode(qml.device("lightning.qubit",wires=17), shots=num_shots, mcm_method="one-shot")
    def _circuit():
        return random_callback(0, prob)

    res=_circuit()
    print(res)
    assert len(res)==num_shots
    percentage1 = len([b for b in res if b])/num_shots
    print(percentage1, prob)
    assert_allclose(percentage1, prob, atol=0.1, rtol=0.1)
2

Hi Sergei, good to hear from you!

I think you’re right that currently the pure_callback will generally work with non-pure functions, I’ve used this fact as well to do some random-number generation without tracking the prng state explicitly. There isn’t anything in the pipeline right now that would break such usage as far as I know, but the idea behind requiring purity in the interface is that it allows you to do additional optimization that you can’t do without (e.g. move these calls around, memoize them, do CSE, etc).

So even though we don’t do any of this right now I can’t guarantee that we won’t in the future. I would say feel free to exploit this as an implementation detail right now, but if you need this as a guarantee longterm I would suggest a feature request on GH.

Cheers :slight_smile:

3

Hehey David, glad to read you too.

I got it, thanks!

I’m working on a proof of concept, so I think I’ll go on with the current behavior for now. Good to know it will not start failing with the current version. (Yet I have created another issue which might be good to catch)

1 Like