tc.set_backend("pytorch")
class QLayer(nn.Module):
 def __init__(self, n_qubits=5, n_layers=2):
 super(QLayer, self).__init__()
 self.n_qubits = n_qubits
 self.n_layers = n_layers
 self.weights = nn.Parameter(torch.rand(n_layers, n_qubits, 3) * 2 * torch.pi)
 self.K = tc.set_backend("tensorflow")
 qpreds_vmap = self.K.vmap(self.qcirc, vectorized_argnums=0)
 self.qpreds_batch = tc.interfaces.torch_interface(qpreds_vmap, jit=True)
 def qcirc(self, inputs, weights): # Angle Embedding followed Entangling Layers and measurement of PauliZ for each qubit
 c = tc.Circuit(self.n_qubits)
 for i in range(self.n_qubits):
 c.rx(i, theta=inputs[i])
 for j in range(self.n_layers):
 for i in range(self.n_qubits):
 c.r(i, theta=weights[j, i, 0], alpha=weights[j, i, 1], phi=weights[j, i, 2])
 for i in range(self.n_qubits):
 c.cnot(i, (i + 1) % self.n_qubits)
 return self.K.stack([self.K.real(c.expectation_ps(z=[i])) for i in range(self.n_qubits)])
 def forward(self, inputs):
 outputs = self.qpreds_batch(inputs, self.weights)
 return outputs

class Network(nn.Module):
 def __init__(self, HL_dim=64, in_dim=1, out_dim=1):
 super(Network, self).__init__()
 self.model = nn.Sequential(
 nn.Linear(in_dim, n_qubits), nn.Tanh(),
 QLayer(n_qubits=n_qubits, n_layers=n_layers), # Quantum layer
 nn.Linear(n_qubits, out_dim))
 def forward(self, t):
 return self.model(t)
 def compute_loss(self, t):
 u = self.model(t)
 u_t = grad(u, t, grad_outputs=torch.ones_like(u), create_graph=True)[0]
 loss_fun = nn.MSELoss() # Default reduction is 'mean'
 residual_1 = u_t - torch.cos(t)
 eqn_loss_1 = loss_fun(residual_1, torch.zeros_like(residual_1))
 device = u.device
 ## IC loss ##
 IC_loss = loss_fun(u[0], torch.tensor(0.0, device=device))
 return eqn_loss_1, IC_loss

model = PINNs_net().to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
scheduler = ExponentialLR(optimizer, gamma=0.9999)
Nt = 32
# domain dimensions
t_i, t_f = 0.0, torch.pi
t = torch.linspace(t_i, t_f, Nt).requires_grad_(True).to(device)
for epoch in range(epochs + 1):
 # Compute various losses
 eqn_loss_1, IC_loss = model.compute_loss(t.view(-1, 1))
 # Compute the total loss
 total_loss = (eqn_loss_1 + IC_loss) # Strongly respect boundary, initial conditions and symmetry
 # Backward pass
 total_loss.backward()
 optimizer.step()
 optimizer.zero_grad() # Reset grads
 scheduler.step()