VQLS value error when implementing U to obtain b

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?

Hi @adilet_akimshe , welcome to the Forum!

Can you please post the full error traceback that you get? Thanks!