import tensorly as tl
import tlquantum as tlq
from tensorly.tt_matrix import TTMatrix
import torch
from torch import randint, rand, arange, complex128
from torch.optim import Adam
import matplotlib.pyplot as plt
from math import factorial, sqrt
import os
import numpy as np
class Metrology(torch.nn.Module):
 def __init__(self,nqubits,L,dt,approx,ncontraq,ncontral,axes=['x'],scale=0.1,dtype=complex128,device=None):
 super().__init__()
 self.nqubits, self.L, self.dt = nqubits, L, dt
 self.approx, self.ncontraq, self.ncontral, self.axes = approx, ncontraq, ncontral, axes
 self.dtype, self.device = dtype, device
 self.scale=scale
 
 self.Js = torch.nn.Parameter(scale*rand([L])).to(device)
 self.ths = torch.nn.Parameter(scale*rand([L,len(axes)])).to(device)
 self.O = torch.tensor([1.0], device=device)
 self.W = torch.nn.Linear(1,1, device=None, bias=False).double()
 self.sz = Flip(self.nqubits,ncontraq,device=device,dtype=dtype)
 
 self.state = make_state(self.nqubits)
 self.state = tlq.qubits_contract(self.state,self.ncontraq)
 self.make_U1()
 self.make_U2()
 self.Op = make_Sy(self.nqubits)
 def concatenate_parameters(self):
 parameters = []
 for param in self.parameters():
 parameters.append(param.view(-1))
 return torch.cat(parameters)
 def get_concatenated_shape(self):
 concatenated_parameters = self.concatenate_parameters()
 return concatenated_parameters.shape
 
 def make_U1(self): 
 self.U1=[]
 for i in range(self.L):
 self.U1+=get_global_perceptron(self.nqubits, self.ncontraq, self.approx, dt=self.dt, device=self.device, J=self.Js[i])
 for a in range(len(self.axes)):
 self.U1+=get_global_rotation(self.nqubits, self.ncontraq, axis=self.axes[a], param=self.ths[i,a])
 return
 
 def make_U2(self):
 self.U2=[self.sz]
 cnt=0
 for l in range(len(self.U1)):
 mod1 = self.U1[l]
 if cnt%(len(self.axes)+1)==0: #We have a perceptron
 for p in mod1.parameters(): J = p.data
 mod2 = get_global_perceptron(self.nqubits,self.ncontraq,self.approx,grad=False,J=J)
 self.U2 += mod2
 else: #We have an rotation. The axis is given by
 a = self.axes[l%(len(self.axes)+1)-1]
 for p in mod1.parameters(): theta = -p.data
 mod2 = get_global_rotation(self.nqubits,self.ncontraq,axis=a,grad=False,param=theta)
 self.U2 += mod2
 cnt+=1
 return
 
 def forward(self, phi): 
 P = perturbation(self.nqubits,phi)
 c = tlq.TTCircuit(self.U1+P+self.U2, self.ncontraq, self.ncontral)
 #return c.forward_expectation_value(self.state, self.Op).real/(self.nqubits/2)
 return self.W(c.forward_expectation_value(self.state, self.Op).real/(self.nqubits/2))
 
 def forward_ket(self, phi):
 def wrapper(params, phi=phi):
 params = params
 P = perturbation(self.nqubits,phi)
 c = tlq.TTCircuit(self.U1+P+self.U2, self.ncontraq, self.ncontral)
 state = c.to_ket(self.state)
 return state
 
 return wrapper

import tensorcircuit as tc
import torch
import numpy as np
import tensorflow as tf
Nqubits=10
Ntrails=30
L = 10
Nepochs=50
batchsize=20
axes=['x']
Phis = torch.tensor(np.linspace(-0.5,0.5,100), dtype=torch.float32)
Phis_train = 0.5*(2*torch.rand(100)-1)
m = Metrology(Nqubits,L,0.05,1,2,2,scale=0.5,axes=axes)
from tensorcircuit.experimental import qng, qng2
tc.set_backend('pytorch')
# Concatenate parameters
params_generator = m.parameters()
params_pytorch = [torch.Tensor(param.detach().numpy()) for param in params_generator]
# Convert each tensor in the list to the desired dtype
# params_pytorch = [param.type(torch.FloatTensor) for param in params_pytorch]
params_pytorch[1] = params_pytorch[1].squeeze()
params_pytorch[2] = params_pytorch[2].squeeze(dim=1)
params_pytorch = torch.cat(params_pytorch, dim=0)
state = m.forward_ket(Phis[0]) # Should be a Callable which returns a Tensor
qfi_fun = qng(state)
qfi_value = qfi_fun(params_pytorch)

{
	"name": "RuntimeError",
	"message": "You are attempting to call Tensor.requires_grad_() (or perhaps using torch.autograd.functional.* APIs) inside of a function being transformed by a functorch transform. This is unsupported, please attempt to use the functorch transforms (e.g. grad, vjp, jacrev, jacfwd, hessian) or call requires_grad_() outside of a function being transformed instead.",
	"stack": "---------------------------------------------------------------------------
RuntimeError Traceback (most recent call last)
File ~..., in qng.<locals>.wrapper(params, **kws)
 100 psi = f(params)
 101 if mode == \"fwd\":
--> 102 jac = backend.jacfwd(f)(params)
 103 else: # \"rev\"
 104 jac = backend.jacrev(f)(params)
File ~..., in ExtendedBackend.jacfwd.<locals>.wrapper(*args, **kws)
 1507 jjs = []
 1508 for argnum in argnums: # type: ignore
-> 1509 jj = self.vmap(jvp1, vectorized_argnums=2)(
 1510 pf,
 1511 args,
 1512 tuple(
 1513 [
 1514 self.reshape(
 1515 self.eye(self.sizen(arg), dtype=arg.dtype),
 1516 [-1] + list(self.shape_tuple(arg)),
 1517 )
 1518 if i == argnum
 1519 else self.reshape(
 1520 self.zeros(
 1521 [self.sizen(arg), self.sizen(arg)], dtype=arg.dtype
 1522 ),
 1523 [-1] + list(self.shape_tuple(arg)),
 1524 )
 1525 for i, arg in enumerate(args)
 1526 ]
 1527 ),
 1528 )
 1529 jj = self.tree_map(
 1530 partial(
 1531 _transform, input_shape=list(self.shape_tuple(args[argnum]))
 1532 ),
 1533 jj,
 1534 )
 1535 jjs.append(jj)
File ~..., in PyTorchBackend.vmap.<locals>.wrapper(*args, **kws)
 609 def wrapper(*args: Any, **kws: Any) -> Tensor:
 610 in_axes = tuple([0 if i in vectorized_argnums else None for i in range(len(args))]) # type: ignore
--> 611 return torchlib.vmap(f, in_axes, 0)(*args, **kws)
File ~..., in vmap.<locals>.wrapped(*args, **kwargs)
 187 def wrapped(*args, **kwargs):
--> 188 return vmap_impl(func, in_dims, out_dims, randomness, chunk_size, *args, **kwargs)
File ~..., in vmap_impl(func, in_dims, out_dims, randomness, chunk_size, *args, **kwargs)
 262 return _chunked_vmap(func, flat_in_dims, chunks_flat_args,
 263 args_spec, out_dims, randomness, **kwargs)
 265 # If chunk_size is not specified.
--> 266 return _flat_vmap(
 267 func, batch_size, flat_in_dims, flat_args, args_spec, out_dims, randomness, **kwargs
 268 )
File ~..., in doesnt_support_saved_tensors_hooks.<locals>.fn(*args, **kwargs)
 35 @functools.wraps(f)
 36 def fn(*args, **kwargs):
 37 with torch.autograd.graph.disable_saved_tensors_hooks(message):
---> 38 return f(*args, **kwargs)
File ~..., in _flat_vmap(func, batch_size, flat_in_dims, flat_args, args_spec, out_dims, randomness, **kwargs)
 377 try:
 378 batched_inputs = _create_batched_inputs(flat_in_dims, flat_args, vmap_level, args_spec)
--> 379 batched_outputs = func(*batched_inputs, **kwargs)
 380 return _unwrap_batched(batched_outputs, out_dims, vmap_level, batch_size, func)
 381 finally:
File ~..., in return_partial.<locals>.wrapper(*args, **kws)
 62 @wraps(f)
 63 def wrapper(*args: Any, **kws: Any) -> Any:
---> 64 r = f(*args, **kws)
 65 nr = [r[ind] for ind in return_argnums] # type: ignore
 66 if one_input:
File ~..., in PyTorchBackend.jvp(self, f, inputs, v)
 594 v = tuple(v)
 595 # for both tf and torch
 596 # behind the scene: https://j-towns.github.io/2017/06/12/A-new-trick.html
 597 # to be investigate whether the overhead issue remains as in
 598 # https://github.com/renmengye/tensorflow-forward-ad/issues/2
--> 599 return torchlib.autograd.functional.jvp(f, inputs, v)
File ~..., in jvp(func, inputs, v, create_graph, strict)
 421 with torch.enable_grad():
 422 is_inputs_tuple, inputs = _as_tuple(inputs, \"inputs\", \"jvp\")
--> 423 inputs = _grad_preprocess(inputs, create_graph=create_graph, need_graph=True)
 425 if v is not None:
 426 _, v = _as_tuple(v, \"v\", \"jvp\")
File ~..., in _grad_preprocess(inputs, create_graph, need_graph)
 83 res.append(inp.clone())
 84 else:
---> 85 res.append(inp.detach().requires_grad_(need_graph))
 86 return tuple(res)
RuntimeError: You are attempting to call Tensor.requires_grad_() (or perhaps using torch.autograd.functional.* APIs) inside of a function being transformed by a functorch transform. This is unsupported, please attempt to use the functorch transforms (e.g. grad, vjp, jacrev, jacfwd, hessian) or call requires_grad_() outside of a function being transformed instead."
}