How to broadcast non-trainable parameters

Hello! I’m currently building a QNode that takes in a list of parameters and data called phi which is just meant to specify something about my model and is therefore not supposed to be trainable.

dev = qml.device('default.qubit', wires=n_wires, shots=None)

def layers(params, L):
    for i in range(L):
        qml.evolve(qml.sum(*[qml.PauliZ(i) @ qml.PauliZ(i+1)  for i in range(n_wires-1)]), coeff = params[0])
        qml.evolve(qml.sum(*[qml.PauliX(i)  for i in range(n_wires)]), coeff = params[1])
        qml.evolve(qml.sum(*[qml.PauliY(i)  for i in range(n_wires)]), coeff = params[2])

@qml.qnode(dev, interface='jax')
def simplistic_circuit(params, phi=jnp.array([0.001]), L=1):

    S_0 = n_wires/2

    for i in range(n_wires): # Making the initial CSS
        qml.Hadamard(wires=i)

    qml.Barrier()

    layers(params, L)

    qml.Barrier()

    for i in range(n_wires): # Perturbation
        qml.RY(phi[0], wires = i)  

    qml.Barrier()

    qml.adjoint(layers)(params, L)

    return qml.expval((phi[0]/(2*S_0))*qml.sum(*[qml.PauliY(i) for i in range(n_wires)]))

params_normal = jnp.array([0.001,0.002,0.003])
phi_broadcast = simplistic_circuit(params_normal, phi = jnp.array([0.001, 0.01]))

print(phi_broadcast)

which outputs

-8.000233331980234e-10

instead of a 1D array with 2 elements as I’m expecting (each element is the expectation value of my observable when the phi value is 0.001 and 0.01 respectively).

Is it possible to do parameter broadcasting for non-trainable parameters. And if so, what am I doing wrong here?

Here is the output of my qml.about():

Name: PennyLane
Version: 0.33.1
Summary: PennyLane is a Python quantum machine learning library by Xanadu Inc.
Home-page: https://github.com/PennyLaneAI/pennylane
Author: 
Author-email: 
License: Apache License 2.0
Location: /Users/.../python3.11/site-packages
Requires: appdirs, autograd, autoray, cachetools, networkx, numpy, pennylane-lightning, requests, rustworkx, scipy, semantic-version, toml, typing-extensions
Required-by: PennyLane-Lightning

Platform info:           macOS-13.4.1-x86_64-i386-64bit
Python version:          3.11.5
Numpy version:           1.26.2
Scipy version:           1.11.4
Installed devices:
- default.gaussian (PennyLane-0.33.1)
- default.mixed (PennyLane-0.33.1)
- default.qubit (PennyLane-0.33.1)
- default.qubit.autograd (PennyLane-0.33.1)
- default.qubit.jax (PennyLane-0.33.1)
- default.qubit.legacy (PennyLane-0.33.1)
- default.qubit.tf (PennyLane-0.33.1)
- default.qubit.torch (PennyLane-0.33.1)
- default.qutrit (PennyLane-0.33.1)
- null.qubit (PennyLane-0.33.1)
- lightning.qubit (PennyLane-Lightning-0.33.1)

Thank you for all your help!

Hey @NickGut0711,

You can most certainly do broadcasting with non-trainable parameters :slight_smile: here’s an example:

dev = qml.device("default.qubit")

@qml.qnode(dev)
def circuit(trainable, non_trainable):
    qml.RX(non_trainable, wires=0)
    qml.RY(trainable, wires=0)
    return qml.expval(qml.PauliZ(0))

trainable = np.array([0.1])
non_trainable = np.array([0.1, 0.2], requires_grad=False)

print(circuit(trainable, non_trainable))
print(qml.jacobian(circuit)(trainable, non_trainable))
[0.99003329 0.97517033]
[[-0.09933467]
 [-0.0978434 ]]

As it pertains to your example, this might be a little problematic:

   for i in range(n_wires): # Perturbation
        qml.RY(phi[0], wires = i)  

When you’re accessing an element of something within a QNode and that something can have a leading dimension, then you need to be careful. phi[0] is going to be a scalar — not something with a leading dimension to broadcast over.

Let me know if that helps!

Hi @isaacdevlugt! Thank you so much :grin:

So, just to make sure. I can leave the argument are just phi and when I call the QNode later on it can either be some scalar or an array with some leading dimension?

Hello @isaacdevlugt. I have another question in regards to this.

I left the variable as phi to do parameter broadcasting in the following code

# Defining the quantum circuit

def U1(params, L, num_trotter_steps = 10): # We didn't include num_trotter_steps and it seemed to work

    # n_js = L
    # n_ws = L
    # n_theta_xs = n_wires * L
    # n_theta_ys = n_wires * L

    start_index = 0
    n_params_in_layer = (2 + 2) * L

    for i in range(L):
        params = params[start_index:start_index + n_params_in_layer]
        H = H_perceptron
        qml.evolve(perceptron.H)(params[:2], t)
        for j in range(n_wires): # Single-qubit X rotations
            qml.RX(params[2], wires=j)
        for k in range(n_wires): # Single-qubit Y rotations
            qml.RY(params[3], wires=k)
        start_index += n_params_in_layer

def get_Sy(n_wires, phi):
    S_0 = n_wires/2
    c = 0

    for i in range(n_wires):
        c += (phi/(2*S_0))*qml.PauliY(wires=i)

    return c

@qml.qnode(dev, interface='jax')
def circuit(params, phi = 0.001, L = 1):

    for i in range(n_wires): # Making the initial CSS
        qml.Hadamard(wires=i)

    U1(params, L)

    qml.Barrier()

    for i in range(n_wires): # Perturbation
        qml.RY(phi, wires = i)

    qml.Barrier()

    qml.adjoint(U1)(params, L)

    return qml.expval(get_Sy(n_wires, phi))

def my_model(params, phi, L):
    return circuit(params, phi = phi, L = L)

def loss_fn(param_vector, targets, phi, L):
    param_list = perceptron.vector_to_hamiltonian_parameters(param_vector)
    predictions = my_model(param_list, phi, L)
    loss = jnp.sum((targets - predictions) ** 2 / len(targets))
    return loss

random_seed = 25

params = perceptron.get_random_parameter_vector(random_seed)

opt = optax.adam(learning_rate=0.3)
opt_state = opt.init(params)

def update_step(params, opt_state, targets, phi, L):
    loss_val, grads = jax.value_and_grad(loss_fn)(params, targets, phi, L)
    updates, opt_state = opt.update(grads, opt_state)
    params = optax.apply_updates(params, updates)
    return params, opt_state, loss_val

phis_train = jnp.linspace(-0.005,0.005,100)

loss_history = []

for i in range(100):
    params, opt_state, loss_val = update_step(params, opt_state, phis_train, phis_train, L)

    if i % 5 == 0:
        print(f"Step: {i} Loss: {loss_val}")

    loss_history.append(loss_val)

Now, when I try to optimize my circuit I get the error:

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
/Users/.../.ipynb Cell 19 line 8
      5 loss_history = []
      7 for i in range(100):
----> 8     params, opt_state, loss_val = update_step(params, opt_state, phis_train, phis_train, L)
     10     if i % 5 == 0:
     11         print(f"Step: {i} Loss: {loss_val}")

/Users/.../.ipynb Cell 19 line 5
      4 def update_step(params, opt_state, targets, phi, L):
----> 5     loss_val, grads = jax.value_and_grad(loss_fn)(params, targets, phi, L)
      6     updates, opt_state = opt.update(grads, opt_state)
      7     params = optax.apply_updates(params, updates)

    [... skipping hidden 8 frame]

/Users/.../.ipynb Cell 19 line 9
      7 def loss_fn(param_vector, targets, phi, L):
      8     param_list = perceptron.vector_to_hamiltonian_parameters(param_vector)
----> 9     predictions = my_model(param_list, phi, L)
     10     loss = jnp.sum((targets - predictions) ** 2 / len(targets))
     11     return loss

/Users/.../.ipynb Cell 19 line 4
      3 def my_model(params, phi, L):
----> 4     return circuit(params, phi = phi, L = L)

File ~/.../pennylane/qnode.py:970, in QNode.__call__(self, *args, **kwargs)
    967         kwargs["shots"] = _get_device_shots(self._original_device)
    969 # construct the tape
--> 970 self.construct(args, kwargs)
    972 cache = self.execute_kwargs.get("cache", False)
    973 using_custom_cache = (
    974     hasattr(cache, "__getitem__")
    975     and hasattr(cache, "__setitem__")
    976     and hasattr(cache, "__delitem__")
    977 )

File ~/.../pennylane/qnode.py:856, in QNode.construct(self, args, kwargs)
    853     self.interface = qml.math.get_interface(*args, *list(kwargs.values()))
    855 with qml.queuing.AnnotatedQueue() as q:
--> 856     self._qfunc_output = self.func(*args, **kwargs)
    858 self._tape = QuantumScript.from_queue(q, shots)
    860 params = self.tape.get_parameters(trainable_only=False)

/Users/.../.ipynb Cell 19 line 2
     25 qml.Barrier()
     27 qml.adjoint(U1)(params, L)
---> 29 return qml.expval(get_Sy(n_wires, phi))

File ~/.../pennylane/measurements/expval.py:61, in expval(op)
     58 if isinstance(op, MeasurementValue):
     59     return ExpectationMP(obs=op)
---> 61 if not op.is_hermitian:
     62     warnings.warn(f"{op.name} might not be hermitian.")
     64 return ExpectationMP(obs=op)

File ~/.../pennylane/ops/op_math/sum.py:176, in Sum.is_hermitian(self)
    173     if not math.is_abstract(coeffs_list[0]):
    174         return not any(math.iscomplex(c) for c in coeffs_list)
--> 176 return all(s.is_hermitian for s in self)

File ~/.../pennylane/ops/op_math/sum.py:176, in <genexpr>(.0)
    173     if not math.is_abstract(coeffs_list[0]):
    174         return not any(math.iscomplex(c) for c in coeffs_list)
--> 176 return all(s.is_hermitian for s in self)

File ~/.../pennylane/ops/op_math/sum.py:176, in Sum.is_hermitian(self)
    173     if not math.is_abstract(coeffs_list[0]):
    174         return not any(math.iscomplex(c) for c in coeffs_list)
--> 176 return all(s.is_hermitian for s in self)

File ~/.../pennylane/ops/op_math/sum.py:176, in <genexpr>(.0)
    173     if not math.is_abstract(coeffs_list[0]):
    174         return not any(math.iscomplex(c) for c in coeffs_list)
--> 176 return all(s.is_hermitian for s in self)

    [... skipping similar frames: <genexpr> at line 176 (1 times), Sum.is_hermitian at line 176 (1 times)]

File ~/.../pennylane/ops/op_math/sum.py:176, in Sum.is_hermitian(self)
    173     if not math.is_abstract(coeffs_list[0]):
    174         return not any(math.iscomplex(c) for c in coeffs_list)
--> 176 return all(s.is_hermitian for s in self)

File ~/.../pennylane/ops/op_math/sum.py:176, in <genexpr>(.0)
    173     if not math.is_abstract(coeffs_list[0]):
    174         return not any(math.iscomplex(c) for c in coeffs_list)
--> 176 return all(s.is_hermitian for s in self)

File ~/.../pennylane/ops/op_math/sprod.py:193, in SProd.is_hermitian(self)
    189 @property
    190 def is_hermitian(self):
    191     """If the base operator is hermitian and the scalar is real,
    192     then the scalar product operator is hermitian."""
--> 193     return self.base.is_hermitian and not qml.math.iscomplex(self.scalar)

File ~/.../jax/_src/array.py:261, in ArrayImpl.__bool__(self)
    258 def __bool__(self):
    259   # deprecated 2023 September 18.
    260   # TODO(jakevdp) change to warn_on_empty=False
--> 261   core.check_bool_conversion(self, warn_on_empty=True)
    262   return bool(self._value)

File ~/.../jax/_src/core.py:625, in check_bool_conversion(arr, warn_on_empty)
    622     raise ValueError("The truth value of an empty array is ambiguous. Use "
    623                      "`array.size > 0` to check that an array is not empty.")
    624 if arr.size > 1:
--> 625   raise ValueError("The truth value of an array with more than one element is "
    626                     "ambiguous. Use a.any() or a.all()")

ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()

This was working fine when phi was just a scalar. But now, since it is an array, it is not working. It has to do with my get_Sy function. Thank you!

Hey @NickGut0711,

I think the problem is in loss_fn

    loss = jnp.sum((targets - predictions) ** 2 / len(targets))

When you add/subtract targets and predictions and the dimensions aren’t the same it can cause an error like this. Are targets and predictions the same dimension?