code: def fidelity(state0, state1, check_state=False, c_dtype="complex128"):
      r"""Compute the fidelity for two states (given as density matrices) acting on quantum
      systems with the same size.
  
      The fidelity for two mixed states given by density matrices :math:`\rho` and :math:`\sigma`
      is defined as
  
      .. math::
          F( \rho , \sigma ) = \text{Tr}( \sqrt{\sqrt{\rho} \sigma \sqrt{\rho}})^2
  
      .. note::
          It supports all interfaces (Numpy, Autograd, Torch, Tensorflow and Jax). The second state is coerced
          to the type and dtype of the first state. The fidelity is returned in the type of the interface of the
          first state.
  
      Args:
          state0 (tensor_like): ``(2**N, 2**N)`` or ``(batch_dim, 2**N, 2**N)`` density matrix.
          state1 (tensor_like): ``(2**N, 2**N)`` or ``(batch_dim, 2**N, 2**N)`` density matrix.
          check_state (bool): If True, the function will check the validity of both states; that is,
              (shape, trace, positive-definitiveness) for density matrices.
          c_dtype (str): Complex floating point precision type.
  
      Returns:
          float: Fidelity between the two quantum states.
  
      **Example**
  
      To find the fidelity between two state vectors, call :func:`~.math.dm_from_state_vector` on the
      inputs first, e.g.:
  
      >>> state0 = qml.math.dm_from_state_vector([0.98753537-0.14925137j, 0.00746879-0.04941796j])
      >>> state1 = qml.math.dm_from_state_vector([0.99500417+0.j, 0.09983342+0.j])
      >>> qml.math.fidelity(state0, state1)
      0.9905158135644924
  
      To find the fidelity between two density matrices, they can be passed directly:
  
      >>> state0 = [[1, 0], [0, 0]]
      >>> state1 = [[0, 0], [0, 1]]
      >>> qml.math.fidelity(state0, state1)
      0.0
  
      .. seealso:: :func:`pennylane.math.fidelity_statevector` and :func:`pennylane.qinfo.transforms.fidelity`
  
      """
      # Cast as a c_dtype array
      state0 = cast(state0, dtype=c_dtype)
      state1 = cast(state1, dtype=c_dtype)
  
      if check_state:
          _check_density_matrix(state0)
          _check_density_matrix(state1)
  
      if qml.math.shape(state0)[-1] != qml.math.shape(state1)[-1]:
          raise qml.QuantumFunctionError(
              "The two states must have the same number of wires."
          )
  
      # Two mixed states
      _register_vjp(state0, state1)
      fid = qml.math.compute_fidelity(state0, state1)
      return fid
  
file_path: pennylane/pennylane/math/fidelity.py
imports: from functools import lru_cache
  import autograd
  import autoray
  import pennylane
  from .utils import cast
  from .quantum import _check_density_matrix, _check_state_vector
label_desc: Use image node_id 2 for calling a global function with example usage: fidelity(state0, state1, check_state, c_dtype) and returns: fid