[BUG] final transforms with non-trainable parameters

[albi3ro]

Expected behavior

I'd expect to get the same answer as when the AmplitudeEmbedding with null behavior is removed.

In an ideal world, a gradient transform would be able to detect that the trainability information is out of date and recalculate it.

Actual behavior

Traceback below.

The device preprocessing has to decompose AmplitudeEmbedding, so it eliminates the trainablility information set at the qnode level. The default behaviour for QuantumScript.trainable_params then assumes the resulting StatePrep should be trainable, even though the original AmplitudeEmbedding wasn't.

Since parameter shift cannot differentiate StatePrep, this causes an error.

Additional information

A similar example is:

@qml.gradients.param_shift
@qml.compile
@qml.qnode(qml.device('default.qubit'))
def circuit(x):
    qml.QubitUnitary(np.eye(2), 0)
    qml.RX(x, wires=0)
    return qml.expval(qml.PauliZ(0))

circuit(qml.numpy.array(0.5))

Source code

@qml.gradients.param_shift
@qml.qnode(qml.device('default.qubit'))
def circuit(x):
    qml.AmplitudeEmbedding([1,0], wires=0)
    qml.RX(x, wires=0)
    return qml.expval(qml.PauliZ(0))

circuit(qml.numpy.array(0.5))

Tracebacks

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[10], line 8
      5     qml.RX(x, wires=0)
      6     return qml.expval(qml.PauliZ(0))
----> 8 circuit(qml.numpy.array(0.5))

File ~/Prog/pennylane/pennylane/workflow/qnode.py:1038, in QNode.__call__(self, *args, **kwargs)
   1033     if hybrid:
   1034         argnums = full_transform_program[-1]._kwargs.pop(
   1035             "argnums", None
   1036         )  # pylint: disable=protected-access
-> 1038         full_transform_program._set_all_classical_jacobians(
   1039             self, args, kwargs, argnums
   1040         )  # pylint: disable=protected-access
   1041         full_transform_program._set_all_argnums(
   1042             self, args, kwargs, argnums
   1043         )  # pylint: disable=protected-access
   1045 # pylint: disable=unexpected-keyword-arg

File ~/Prog/pennylane/pennylane/transforms/core/transform_program.py:420, in TransformProgram._set_all_classical_jacobians(self, qnode, args, kwargs, argnums)
    416     raise qml.QuantumFunctionError(
    417         "argnum does not work with the Jax interface. You should use argnums instead."
    418     )
    419 sub_program = TransformProgram(self[0:index])
--> 420 classical_jacobian = jacobian(
    421     classical_preprocessing, sub_program, argnums, *args, **kwargs
    422 )
    423 qnode.construct(args, kwargs)
    424 tapes, _ = sub_program((qnode.tape,))

File ~/Prog/pennylane/pennylane/transforms/core/transform_program.py:376, in TransformProgram._set_all_classical_jacobians.<locals>.jacobian(classical_function, program, argnums, *args, **kwargs)
    373 classical_function = partial(classical_function, program)
    375 if qnode.interface == "autograd":
--> 376     jac = qml.jacobian(classical_function, argnum=argnums)(*args, **kwargs)
    378 if qnode.interface == "tf":
    379     import tensorflow as tf  # pylint: disable=import-outside-toplevel

File ~/Prog/pennylane/pennylane/_grad.py:455, in jacobian.<locals>._jacobian_function(*args, **kwargs)
    449 if not _argnum:
    450     warnings.warn(
    451         "Attempted to differentiate a function with no trainable parameters. "
    452         "If this is unintended, please add trainable parameters via the "
    453         "'requires_grad' attribute or 'argnum' keyword."
    454     )
--> 455 jac = tuple(_jacobian(func, arg)(*args, **kwargs) for arg in _argnum)
    457 return jac[0] if unpack else jac

File ~/Prog/pennylane/pennylane/_grad.py:455, in <genexpr>(.0)
    449 if not _argnum:
    450     warnings.warn(
    451         "Attempted to differentiate a function with no trainable parameters. "
    452         "If this is unintended, please add trainable parameters via the "
    453         "'requires_grad' attribute or 'argnum' keyword."
    454     )
--> 455 jac = tuple(_jacobian(func, arg)(*args, **kwargs) for arg in _argnum)
    457 return jac[0] if unpack else jac

File ~/Prog/pl311/lib/python3.11/site-packages/autograd/wrap_util.py:20, in unary_to_nary.<locals>.nary_operator.<locals>.nary_f(*args, **kwargs)
     18 else:
     19     x = tuple(args[i] for i in argnum)
---> 20 return unary_operator(unary_f, x, *nary_op_args, **nary_op_kwargs)

File ~/Prog/pl311/lib/python3.11/site-packages/autograd/differential_operators.py:60, in jacobian(fun, x)
     50 @unary_to_nary
     51 def jacobian(fun, x):
     52     """
     53     Returns a function which computes the Jacobian of `fun` with respect to
     54     positional argument number `argnum`, which must be a scalar or array. Unlike
   (...)
     58     (out1, out2, ...) then the Jacobian has shape (out1, out2, ..., in1, in2, ...).
     59     """
---> 60     vjp, ans = _make_vjp(fun, x)
     61     ans_vspace = vspace(ans)
     62     jacobian_shape = ans_vspace.shape + vspace(x).shape

File ~/Prog/pl311/lib/python3.11/site-packages/autograd/core.py:10, in make_vjp(fun, x)
      8 def make_vjp(fun, x):
      9     start_node = VJPNode.new_root()
---> 10     end_value, end_node =  trace(start_node, fun, x)
     11     if end_node is None:
     12         def vjp(g): return vspace(x).zeros()

File ~/Prog/pl311/lib/python3.11/site-packages/autograd/tracer.py:10, in trace(start_node, fun, x)
      8 with trace_stack.new_trace() as t:
      9     start_box = new_box(x, t, start_node)
---> 10     end_box = fun(start_box)
     11     if isbox(end_box) and end_box._trace == start_box._trace:
     12         return end_box._value, end_box._node

File ~/Prog/pl311/lib/python3.11/site-packages/autograd/wrap_util.py:15, in unary_to_nary.<locals>.nary_operator.<locals>.nary_f.<locals>.unary_f(x)
     13 else:
     14     subargs = subvals(args, zip(argnum, x))
---> 15 return fun(*subargs, **kwargs)

File ~/Prog/pennylane/pennylane/transforms/core/transform_program.py:356, in TransformProgram._set_all_classical_jacobians.<locals>.classical_preprocessing(program, *args, **kwargs)
    354 tape = qnode.qtape
    355 tapes, _ = program((tape,))
--> 356 res = tuple(qml.math.stack(tape.get_parameters(trainable_only=True)) for tape in tapes)
    357 if len(tapes) == 1:
    358     return res[0]

File ~/Prog/pennylane/pennylane/transforms/core/transform_program.py:356, in <genexpr>(.0)
    354 tape = qnode.qtape
    355 tapes, _ = program((tape,))
--> 356 res = tuple(qml.math.stack(tape.get_parameters(trainable_only=True)) for tape in tapes)
    357 if len(tapes) == 1:
    358     return res[0]

File ~/Prog/pennylane/pennylane/math/multi_dispatch.py:151, in multi_dispatch.<locals>.decorator.<locals>.wrapper(*args, **kwargs)
    148 interface = interface or get_interface(*dispatch_args)
    149 kwargs["like"] = interface
--> 151 return fn(*args, **kwargs)

File ~/Prog/pennylane/pennylane/math/multi_dispatch.py:488, in stack(values, axis, like)
    459 """Stack a sequence of tensors along the specified axis.
    460 
    461 .. warning::
   (...)
    485        [5.00e+00, 8.00e+00, 1.01e+02]], dtype=float32)>
    486 """
    487 values = np.coerce(values, like=like)
--> 488 return np.stack(values, axis=axis, like=like)

File ~/Prog/pl311/lib/python3.11/site-packages/autoray/autoray.py:80, in do(fn, like, *args, **kwargs)
     31 """Do function named ``fn`` on ``(*args, **kwargs)``, peforming single
     32 dispatch to retrieve ``fn`` based on whichever library defines the class of
     33 the ``args[0]``, or the ``like`` keyword argument if specified.
   (...)
     77     <tf.Tensor: id=91, shape=(3, 3), dtype=float32>
     78 """
     79 backend = choose_backend(fn, *args, like=like, **kwargs)
---> 80 return get_lib_fn(backend, fn)(*args, **kwargs)

File ~/Prog/pennylane/pennylane/numpy/wrapper.py:117, in tensor_wrapper.<locals>._wrapped(*args, **kwargs)
    114         tensor_kwargs["requires_grad"] = _np.any([i.requires_grad for i in tensor_args])
    116 # evaluate the original object
--> 117 res = obj(*args, **kwargs)
    119 if isinstance(res, _np.ndarray):
    120     # only if the output of the object is a ndarray,
    121     # then convert to a PennyLane tensor
    122     res = tensor(res, **tensor_kwargs)

File ~/Prog/pl311/lib/python3.11/site-packages/autograd/numpy/numpy_wrapper.py:94, in stack(arrays, axis)
     92 shapes = set(arr.shape for arr in arrays)
     93 if len(shapes) != 1:
---> 94     raise ValueError('all input arrays must have the same shape')
     96 result_ndim = arrays[0].ndim + 1
     97 if not -result_ndim <= axis < result_ndim:

ValueError: all input arrays must have the same shape

System information

Master

Existing GitHub issues

• [X] I have searched existing GitHub issues to make sure the issue does not already exist.

[AnuravModak]

please let me know if i can help in this in any way.

[albi3ro]

@AnuravModak Thanks for the offer.

There may be a simpler, patch fix to this problem, but in my mind this is a symptom of a deeper structural issue, that would require a deeper structural fix.

Basically gradients take QuantumScript.trainable_params as the source of truth as what is trainable or not, and that property defaults to "everything is trainable" if it hasn't been explicitly set.

The problem is that it's extremely hard to keep track of the how trainable_params might get updated throughout transforms, so we just don't track how it gets transformed. compile might eliminate parameters. decompose and expansions might break one parameter into multiple. Tracking the indices through that process would add way too much complexity. So currently we just don't. The trainable_params are set in the qnode and right before we hit the ML boundary.

We could change the default for trainable params to "check what's trainable", but that would be a breaking change that would require buy-in from a variety of different people.