Incompatible shape in solve_normal_cg about jaxopt HOT 6 CLOSED

def _rmatvec(matvec, x):
  """Computes A^T x, from matvec(x) = A x, where A is square."""
  transpose = jax.linear_transpose(matvec, x)
  return transpose(x)[0]

def _rmatvec(matvec, x): # x is of size N 
  transpose = jax.linear_transpose(matvec, dummy_vector_of_size_P)
  return transpose(x)[0]

def solve_normal_cg(matvec: Callable,
                    b: Any,
                    ridge: Optional[float] = None,
                    init: Optional[Any] = None,
                    dummy: Optional[Any] = None,
                    **kwargs) -> Any:

  def _matvec(x):
    """Computes A^T A x."""
    return _normal_matvec(matvec, x)

  if ridge is not None:
    _matvec = _make_ridge_matvec(_matvec, ridge=ridge)

  Ab = _rmatvec(matvec, b, dummy=dummy)

  return jax.scipy.sparse.linalg.cg(_matvec, Ab, x0=init, **kwargs)[0]

if init is None:
  try:
    Ab = _rmatvec(matvec, b)
  except TypeError, e:
    raise TypeError("`init` is compulsory when `matvec` is nonsquare. Original error message:", e)
else:
  Ab = _rmatvec(matvec, b, init)

def solve_normal_cg(matvec: Callable,
                    b: Any,
                    ridge: Optional[float] = None,
                    init: Optional[Any] = None,
                    apply_right: bool = True,
                    **kwargs) -> Any:

  def _matvec(x):
    """Computes A^T A x."""
    return _normal_matvec(matvec, x)

  if ridge is not None:
    _matvec = _make_ridge_matvec(_matvec, ridge=ridge)

  if apply_right: 
    y = _rmatvec(matvec, b)
  else:
    y = b

  return jax.scipy.sparse.linalg.cg(_matvec, y, x0=init, **kwargs)[0]

def _rmatvec(matvec, x, example):
  """Computes A^T x, from matvec(z) = A z, given an example input z for matvec"""
  transpose = jax.linear_transpose(matvec, example)
  return transpose(x)[0]

def solve_normal_cg(matvec: Callable,
                    b: Any,
                    ridge: Optional[float] = None,
                    init: Optional[Any] = None,
                    example_solution: Optional[Any] = None,
                    **kwargs) -> Any:
  """
  [...]

  For example_solution, any object with shape/dtype attributes suffices.
  Consider using jax.eval_shape if that's helpful.
  """

  def _matvec(x):
    """Computes A^T A x."""
    return _normal_matvec(matvec, x)

  if ridge is not None:
    _matvec = _make_ridge_matvec(_matvec, ridge=ridge)

  if example_solution is None:
    example_solution = init
  if example_solution is None:
    example_solution = b  # assuming square matvec

  Ab = _rmatvec(matvec, b, example_solution)

  return jax.scipy.sparse.linalg.cg(_matvec, Ab, x0=init, **kwargs)[0]