code: def purity(rho: Density) -> bk.BKTensor:
      
      tensor = rho.tensor
      N = rho.qubit_nb
      matrix = bk.reshape(tensor, [2**N, 2**N])
      return bk.trace(bk.matmul(matrix, matrix))
docs: Calculate the purity of a mixed quantum state.
  
      Purity, defined as tr(rho^2), has an upper bound of 1 for a pure state,
      and a lower bound of 1/D (where D is the Hilbert space dimension) for a
      competently mixed state.
  
      Two closely related measures are the linear entropy, 1- purity, and the
      participation ratio, 1/purity.