I am trying to construct a density matrix of shape 256x256 from a set of T matrices. These T matrices are all Cholesky-decomposed matrices. But I am not sure if the code is right or my approach is right; one of the reasons is that when I change from tf.einsum to tf.matmul some error is raised. I am not very well used to use einsum, that is the main reason I am trying to change the following code:
def density_matrix_from_T(tmatrix):
"""
Gets density matrices from T matrices and normalizes them.
Args:
tmatrix (`tf.Tensor`): A tensor (N, hilbert_size, hilbert_size)
representing N valid T matrices.
Returns:
rho (`tf.Tensor`): A tensor of shape (N, hilbert_size, hilbert_size)
representing N density matrices.
"""
T = tmatrix
T_dagger = tf.transpose(T, perm=[0, 2, 1], conjugate=True)
proper_dm = tf.matmul(T_dagger, T)
all_traces = tf.linalg.trace(proper_dm)
all_traces = tf.reshape(1 / all_traces, (-1, 1))
rho = tf.einsum("bij,bk->bij", proper_dm, all_traces)
return rho
where the T matrix came from:
def clean_cholesky(img):
"""
Cleans an input matrix to make it the Cholesky decomposition matrix T
Args:
img (`tf.Tensor`): a tensor of shape (batch_size, hilbert_size,
hilbert_size, 2)
representing random outputs from a neural netowrk.
The last dimension is for separating the real and
imaginary part.
Returns:
T (`tf.Tensor`): a 3D tensor (N, hilbert_size, hilbert_size)
representing N matrices used for Cholesky decomp.
"""
real = img[:, :, :, 0]
imag = img[:, :, :, 1]
diag_all = tf.linalg.diag_part(imag, k=0, padding_value=0)
diags = tf.linalg.diag(diag_all)
imag = imag - diags
imag = tf.linalg.band_part(imag, -1, 0)
real = tf.linalg.band_part(real, -1, 0)
T = tf.complex(real, imag)
return T
When I use it like this, it works fine but when I compute Fidelity to compare the reconstructed density matrix with a given original matrix, it gives a very poor result (something like 0.01 when a good result is close to 1). Anyway, I believe the problem is with rho = tf.einsum("bij,bk->bij", proper_dm, all_traces) and then I tried to change it to rho = tf.matmul(proper_dm, all_traces) but it raises an error on incompatibility:
InvalidArgumentError: Exception encountered when calling layer 'density_matrix_1' (type DensityMatrix).
{{function_node __wrapped__BatchMatMulV2_device_/job:localhost/replica:0/task:0
device:CPU:0}} Matrix size-incompatible: In[0]: [1,256,256], In[1]: [1,1] [Op:BatchMatMulV2]
Call arguments received by layer 'density_matrix_1' (type DensityMatrix):
• inputs=tf.Tensor(shape=(1, 256, 256, 2), dtype=float64)
• training=True
The arguments bij, bk -> bij are independent of the input size? I mean, does it matter if the input is a matrix of shape 5x5 or 459x459? Or it is just dependent on the shape of the array in the sense that if it is a 2D or 3D array?
If I don't want (and I really don't want) to use tf.einsum how can I correctly change it to use matmul?
