I was trying to implement my own b in the demo provided by the link below. I create the same b that was used in the demo for VQLS, though contstruct it in a slightly different way. The unitary I create is indeed unitary and of correct size (8x8). If I could do this toy problem with the same b as in the demo, I plan to do it with arbitrary b.

https://pennylane.ai/qml/demos/tutorial_vqls/

```
b = b = np.ones(8) / np.sqrt(8)
def prepare_U(): # householder transformation
zero_state_for_circuit = np.zeros(8)
zero_state_for_circuit[0] = 1
v = zero_state_for_circuit - b
v = v/np.linalg.norm(v)
U = np.identity(np.size(v)) - 2 * np.outer(v, v)
return U
# Run the circuit once to prepare the state
initial_state = np.array(prepare_U(), dtype=complex)
def U_b():
"""Unitary matrix rotating the ground state to the problem vector |b> = U_b |0>."""
# for idx in range(n_qubits): # Original code in the demo
# qml.Hadamard(wires=idx)
qml.QubitUnitary(initial_state, wires=range(n_qubits), unitary_check=True) # I want to implement my own circuit
```

But when it comes to optimization, I get a value error “ValueError: all input arrays must have the same shape”.

```
cost_history = []
for it in range(steps):
w, cost = opt.step_and_cost(cost_loc, w) # Value error
print("Step {:3d} Cost_L = {:9.7f}".format(it, cost))
cost_history.append(cost)
```

Can somebody please explain how to fix that?