def state_vector_call(self, state):
state = K.reshape(state, self.tensor_shape)
substate = K.gather_nd(K.transpose(state, self.order), self.result)
norm = K.sum(K.square(K.abs(substate)))
state = substate / K.cast(K.sqrt(norm), dtype=state.dtype)
state = self._append_zeros(state, self.sorted_qubits, self.result)
return K.reshape(state, self.flat_shape)
def density_matrix_call(self, state):
state = K.reshape(state, self.tensor_shape)
substate = K.gather_nd(K.transpose(state, self.order),
self.density_matrix_result)
n = 2**(len(tuple(substate.shape)) // 2)
norm = K.trace(K.reshape(substate, (n, n)))
state = substate / norm
state = self._append_zeros(state, self.sorted_qubits,
self.density_matrix_result)
return K.reshape(state, self.flat_shape)
