KeyError shape,stack,cos on pennylane quantum circuit

[msaroufim]

@anijain2305 would you rather I create an "interesting model" tracker or should I keep creating unique issues for each kind of model

Repro

python -m pip install pennylane

import torch
import pennylane as qml
import torchdynamo

dev = qml.device('default.qubit', wires=2)

@qml.qnode(dev, interface='torch')
def circuit4(phi, theta):
    qml.RX(phi[0], wires=0)
    qml.RZ(phi[1], wires=1)
    qml.CNOT(wires=[0, 1])
    qml.RX(theta, wires=0)
    return qml.expval(qml.PauliZ(0))

def cost(phi, theta):
    return torch.abs(circuit4(phi, theta) - 0.5)**2

phi = torch.tensor([0.011, 0.012], requires_grad=True)
theta = torch.tensor(0.05, requires_grad=True)

opt = torch.optim.Adam([phi, theta], lr = 0.1)

steps = 200

def closure():
    opt.zero_grad()
    loss = cost(phi, theta)
    loss.backward()
    return loss

with torchdynamo.optimize("eager"):
  for i in range(steps):
        opt.step(closure)

Logs

https://gist.github.com/msaroufim/ce9ec004536e762fb5c94eb3ab2670f1

cc @ezyang @wconstab @bdhirsh @anijain2305 @zou3519 @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @aakhundov @kadeng @soumith @ngimel

[ydwu4]

Not sure if the error is caused by pacakge upgrade in pennylane. But here is the repro:

import torch
import pennylane as qml

dev = qml.device('default.qubit', wires=2)

@qml.qnode(dev, interface='torch')
def circuit4(phi, theta):
    qml.RX(phi[0], wires=0)
    qml.RZ(phi[1], wires=1)
    qml.CNOT(wires=[0, 1])
    qml.RX(theta, wires=0)
    return qml.expval(qml.PauliZ(0))

def cost(phi, theta):
    return torch.abs(circuit4(phi, theta) - 0.5)**2

phi = torch.tensor([0.011, 0.012], requires_grad=True)
theta = torch.tensor(0.05, requires_grad=True)

opt = torch.optim.Adam([phi, theta], lr = 0.1)

steps = 200

def closure():
    opt.zero_grad()
    loss = cost(phi, theta)
    loss.backward()
    return loss

def f():
  for i in range(steps):
        opt.step(closure)

torch.compile(f, backend="eager")()

Error log:

Traceback (most recent call last):
  File "/home/yidi/local/pytorch/repro6.py", line 34, in <module>
    torch.compile(f, backend="eager")()
  File "/home/yidi/local/pytorch/torch/_dynamo/eval_frame.py", line 488, in _fn
    return fn(*args, **kwargs)
  File "/home/yidi/local/pytorch/repro6.py", line 32, in f
    opt.step(closure)
  File "/home/yidi/local/pytorch/torch/optim/optimizer.py", line 76, in _use_grad
    ret = func(self, *args, **kwargs)
  File "/home/yidi/local/pytorch/torch/optim/adam.py", line 146, in step
    loss = closure()
  File "/home/yidi/local/pytorch/repro6.py", line 26, in closure
    loss = cost(phi, theta)
  File "/home/yidi/local/pytorch/repro6.py", line 15, in cost
    return torch.abs(circuit4(phi, theta) - 0.5)**2
  File "/home/yidi/local/miniconda3/envs/pytorch-3.10/lib/python3.10/site-packages/pennylane/qnode.py", line 970, in __call__
    self.construct(args, kwargs)
  File "/home/yidi/local/miniconda3/envs/pytorch-3.10/lib/python3.10/site-packages/pennylane/qnode.py", line 856, in construct
    self._qfunc_output = self.func(*args, **kwargs)
  File "/home/yidi/local/pytorch/repro6.py", line 8, in circuit4
    qml.RX(phi[0], wires=0)
  File "/home/yidi/local/pytorch/torch/_dynamo/eval_frame.py", line 654, in catch_errors
    return callback(frame, cache_entry, hooks, frame_state)
  File "/home/yidi/local/pytorch/torch/_dynamo/convert_frame.py", line 721, in _convert_frame
    result = inner_convert(frame, cache_entry, hooks, frame_state)
  File "/home/yidi/local/pytorch/torch/_dynamo/convert_frame.py", line 383, in _convert_frame_assert
    compiled_product = _compile(
  File "/home/yidi/local/pytorch/torch/_dynamo/convert_frame.py", line 664, in _compile
    raise InternalTorchDynamoError(str(e)).with_traceback(
  File "/home/yidi/local/pytorch/torch/_dynamo/convert_frame.py", line 645, in _compile
    guarded_code = compile_inner(code, one_graph, hooks, transform)
  File "/home/yidi/local/pytorch/torch/_dynamo/utils.py", line 244, in time_wrapper
    r = func(*args, **kwargs)
  File "/home/yidi/local/pytorch/torch/_dynamo/convert_frame.py", line 562, in compile_inner
    out_code = transform_code_object(code, transform)
  File "/home/yidi/local/pytorch/torch/_dynamo/bytecode_transformation.py", line 1033, in transform_code_object
    transformations(instructions, code_options)
  File "/home/yidi/local/pytorch/torch/_dynamo/convert_frame.py", line 151, in _fn
    return fn(*args, **kwargs)
  File "/home/yidi/local/pytorch/torch/_dynamo/convert_frame.py", line 527, in transform
    tracer.run()
  File "/home/yidi/local/pytorch/torch/_dynamo/symbolic_convert.py", line 2123, in run
    super().run()
  File "/home/yidi/local/pytorch/torch/_dynamo/symbolic_convert.py", line 818, in run
    and self.step()
  File "/home/yidi/local/pytorch/torch/_dynamo/symbolic_convert.py", line 781, in step
    getattr(self, inst.opname)(inst)
  File "/home/yidi/local/pytorch/torch/_dynamo/symbolic_convert.py", line 470, in wrapper
    return inner_fn(self, inst)
  File "/home/yidi/local/pytorch/torch/_dynamo/symbolic_convert.py", line 1213, in CALL_FUNCTION
    self.call_function(fn, args, {})
  File "/home/yidi/local/pytorch/torch/_dynamo/symbolic_convert.py", line 652, in call_function
    self.push(fn.call_function(self, args, kwargs))
  File "/home/yidi/local/pytorch/torch/_dynamo/variables/builtin.py", line 479, in call_function
    args = [v.realize() for v in args]
  File "/home/yidi/local/pytorch/torch/_dynamo/variables/builtin.py", line 479, in <listcomp>
    args = [v.realize() for v in args]
  File "/home/yidi/local/pytorch/torch/_dynamo/variables/lazy.py", line 56, in realize
    self._cache.realize(self.parents_tracker)
  File "/home/yidi/local/pytorch/torch/_dynamo/variables/lazy.py", line 22, in realize
    self.vt = VariableBuilder(tx, self.source)(self.value)
  File "/home/yidi/local/pytorch/torch/_dynamo/variables/builder.py", line 239, in __call__
    vt = self._wrap(value).clone(**self.options())
  File "/home/yidi/local/pytorch/torch/_dynamo/variables/builder.py", line 693, in _wrap
    elif trace_rules.lookup(value) is not None:
  File "/home/yidi/local/pytorch/torch/_dynamo/trace_rules.py", line 180, in lookup
    if not hashable(obj):
  File "/home/yidi/local/pytorch/torch/_dynamo/utils.py", line 420, in hashable
    hash(x)
  File "/home/yidi/local/miniconda3/envs/pytorch-3.10/lib/python3.10/site-packages/pennylane/operation.py", line 726, in __hash__
    return self.hash
  File "/home/yidi/local/miniconda3/envs/pytorch-3.10/lib/python3.10/site-packages/pennylane/operation.py", line 715, in hash
    str(self.name),
  File "/home/yidi/local/miniconda3/envs/pytorch-3.10/lib/python3.10/site-packages/pennylane/operation.py", line 901, in name
    return self._name
torch._dynamo.exc.InternalTorchDynamoError: 'RX' object has no attribute '_name'

from user code:
   File "/home/yidi/local/miniconda3/envs/pytorch-3.10/lib/python3.10/site-packages/pennylane/ops/qubit/parametric_ops_single_qubit.py", line 75, in __init__
    super().__init__(phi, wires=wires, id=id)

Set TORCH_LOGS="+dynamo" and TORCHDYNAMO_VERBOSE=1 for more information


You can suppress this exception and fall back to eager by setting:
    import torch._dynamo
    torch._dynamo.config.suppress_errors = True

[ezyang]

The problem is we're hashing a user defined object, which can result in arbitrary code execution. We should not do this

  File "/home/yidi/local/pytorch/torch/_dynamo/trace_rules.py", line 180, in lookup
    if not hashable(obj):