class_docstr: Class for the Bures-Wasserstein metric on the HPD manifold.
docstr: Compute the Bures-Wasserstein squared distance.
  
  Compute the Riemannian squared distance between point_a and point_b.
  
  Parameters
  ----------
  point_a : array-like, shape=[..., n, n]
      Point.
  point_b : array-like, shape=[..., n, n]
      Point.
  
  Returns
  -------
  squared_dist : array-like, shape=[...]
      Riemannian squared distance.
docstr_tok: ["Compute","the","Bures-Wasserstein","squared","distance",".","Compute","the","Riemannian","squared","distance","between","point_a","and","point_b",".","Parameters","--","--","--","--","--","point_a",":","array-like",",","shape=","[","...",",","n",",","n","]","Point",".","point_b",":","array-like",",","shape=","[","...",",","n",",","n","]","Point",".","Returns","--","--","--","-","squared_dist",":","array-like",",","shape=","[","...","]","Riemannian","squared","distance","."]
code: def squared_dist(self, point_a, point_b):
      """Compute the Bures-Wasserstein squared distance.
  
      Compute the Riemannian squared distance between point_a and point_b.
  
      Parameters
      ----------
      point_a : array-like, shape=[..., n, n]
          Point.
      point_b : array-like, shape=[..., n, n]
          Point.
  
      Returns
      -------
      squared_dist : array-like, shape=[...]
          Riemannian squared distance.
      """
      product = gs.matmul(point_a, point_b)
      sqrt_product = gs.linalg.sqrtm(product)
      trace_a = gs.trace(point_a)
      trace_b = gs.trace(point_b)
      trace_prod = gs.trace(sqrt_product)
  
      squared_dist = gs.real(trace_a + trace_b - 2.0 * trace_prod)
  
      return gs.where(squared_dist < 0.0, 0.0, squared_dist)
  
imports: import math
  import geomstats.backend
  from geomstats.geometry.base import ComplexVectorSpaceOpenSet
  from geomstats.geometry.complex_matrices import ComplexMatrices, ComplexMatricesMetric
  from geomstats.geometry.complex_riemannian_metric import ComplexRiemannianMetric
  from geomstats.geometry.general_linear import GeneralLinear
  from geomstats.geometry.hermitian_matrices import HermitianMatrices, expmh, powermh
  from geomstats.geometry.matrices import Matrices
  from geomstats.geometry.pullback_metric import PullbackDiffeoMetric
  from geomstats.geometry.spd_matrices import SymMatrixLog, logmh
  from geomstats.integrator import integrate
  from geomstats.vectorization import repeat_out
label_desc: Use image node_id 4 for calling the HPDBuresWassersteinMetric obj's underlying member method code with example usage: obj.squared_dist(point_a, point_b) and returns: gs