code: def distance_wasserstein(A, B, squared=False):
      r"""Wasserstein distance between SPSD/HPSD matrices.
  
      The Wasserstein distance between two SPSD/HPSD matrices :math:`\mathbf{A}`
      and :math:`\mathbf{B}` is [1]_ [2]_:
  
      .. math::
          d(\mathbf{A},\mathbf{B}) =
          \sqrt{ \text{tr} \left(\mathbf{A} + \mathbf{B}
          - 2(\mathbf{B}^{1/2} \mathbf{A} \mathbf{B}^{1/2})^{1/2} \right) }
  
      Parameters
      ----------
      A : ndarray, shape (..., n, n)
          First SPSD/HPSD matrices, at least 2D ndarray.
      B : ndarray, shape (..., n, n)
          Second SPSD/HPSD matrices, same dimensions as A.
      squared : bool, default False
          Return squared distance.
  
          .. versionadded:: 0.5
  
      Returns
      -------
      d : float or ndarray, shape (...,)
          Wasserstein distance between A and B.
  
      See Also
      --------
      distance
  
      References
      ----------
      .. [1] `Optimal transport: old and new
          <https://link.springer.com/book/10.1007/978-3-540-71050-9>`_
          C. Villani. Springer Science & Business Media, 2008, vol. 338
      .. [2] `An extension of Kakutani's theorem on infinite product measures to
          the tensor product of semifinite w*-algebras
          <https://www.ams.org/journals/tran/1969-135-00/S0002-9947-1969-0236719-2/S0002-9947-1969-0236719-2.pdf>`_
          D. Bures. Trans Am Math Soc, 1969, 135, pp. 199-212
      """  # noqa
      _check_inputs(A, B)
      B12 = sqrtm(B)
      d2 = np.trace(
          A + B - 2 * sqrtm(B12 @ A @ B12), axis1=-2, axis2=-1
      )
      d2 = np.maximum(0, d2.real)
      return d2 if squared else np.sqrt(d2)
  
imports: import numpy
  from scipy.linalg import eigvalsh, solve
  from sklearn.metrics import euclidean_distances
  from .base import logm, sqrtm, invsqrtm
label_desc: Use image node_id 11 for calling a global function with example usage: distance_wasserstein(A, B, squared) and returns: unknown