PennyLaneAI/catalyst

Improve handling of execution configs, especially as default arguments to other functions

joeycarter opened this issue · 6 comments

We noticed in #1117 that using qml.devices.DefaultExecutionConfig, which is a mutable, global object, as default arguments to other functions caused a test to fail in some contexts, most likely because of a race condition where some other function from another test had modified this global variable.

This guide on common Python gotchas explains the fundamental issue, and recommends to generally avoid patterns like

def append_to(element, to=[]):
    to.append(element)
    return to

where a (mutable) empty list is provided as the default argument to the to input argument, and prefer the following pattern instead:

def append_to(element, to=None):
    if to is None:
        to = []
    to.append(element)
    return to

This search query shows where qml.devices.DefaultExecutionConfig is used as a default argument in the Catalyst code base. We should refactor this code to avoid the mutable-default-argument gotcha.

Note on the somewhat unexpected behavior of mutable default arguments:

def f(x=[]):
   x.append(0)
   return x
 
>>> f()
[0]
>>> f()  #  a "sane" person would expect [0]
[0, 0]
>>> f()  #  a "sane" person would expect [0]
[0, 0, 0]

I spoke with @dime10 about this. For this issue, we'll fix up where we modify the global DefaultExecutionConfig, e.g. in jax_tracer.py, treating it as if it were immutable. If an attribute of the class needs to modified, we'll use dataclass.replace, which creates a new object of the same type, replacing fields with the new values. We will also revert #1117 so that the tests use the original call signatures.

Note that there is also a change proposed in PennyLane that makes DefaultExecutionConfig immutable: PennyLaneAI/pennylane#6245.

Small implementation detail: I think we'll need to use deepcopy on DefaultExecutionConfig, rather than dataclass.replace, which performs a shallow copy. This is because several of the config parameters are stored as dictionaries, which should also not be modified in the global object.

Small implementation detail: I think we'll need to use deepcopy on DefaultExecutionConfig, rather than dataclass.replace, which performs a shallow copy. This is because several of the config parameters are stored as dictionaries, which should also not be modified in the global object.

Hmm, we might want to bring in @albi3ro for this question, but replace is the strategy they use in PennyLane. I think deepcopy would be an issue once the config is immutable because the copy will still be immutable.

I actually didn't even realize that replace was a shallow copy, so that is even more of an issue than I was originally thinking.

Summarizing some discussions, I think our path forward should be:

  1. Use replace to update properties, in addition to gradient_keyword_args and device_options.
  2. Defaulting to execution_config: Optional[ExecutionConfig] = None, instead of DefaultExecutionConfig.
  3. Making ExecutionConfig freezable and/or frozen.

Thanks for finding all these things, and I'm sorry for putting this design in to begin with. Hopefully we can fix these mistakes in the next few months.

Thanks @albi3ro! I agree, I think this is the way forward.