Fix: LPU diag argument not passed

Author: heptaflar

uncertaintyx/interface/core.py

@@ -76,7 +76,7 @@ def lpu(
         :param diag: To return only the diagonal elements of :math:`U(Y)`.
         :returns: :math:`U(Y) \in \mathbb{R}^{M \times n \times n}`.
         """
-        return lpu_x(x.ndim - 1, self.jac(x), u)
+        return lpu_x(x.ndim - 1, self.jac(x), u, diag)
 
     @property
     @abstractmethod
@@ -159,7 +159,7 @@ def lpu_p(
         :param diag: To return only the diagonal elements of :math:`U(Y)`.
         :returns: :math:`U(Y) \in \mathbb{R}^{M \times n \times n}`.
         """
-        return lpu_p(p.ndim, self.jac_p(p, x), u)
+        return lpu_p(p.ndim, self.jac_p(p, x), u, diag)
 
     def lpu_x(
         self, p: np.ndarray, x: np.ndarray, u: np.ndarray, diag: bool = False
@@ -190,7 +190,7 @@ def lpu_x(
         :param diag: To return only the diagonal elements of :math:`U(Y)`.
         :returns: :math:`U(Y) \in \mathbb{R}^{M \times n \times n}`.
         """
-        return lpu_x(x.ndim - 1, self.jac_x(p, x), u)
+        return lpu_x(x.ndim - 1, self.jac_x(p, x), u, diag)
 
     @abstractmethod
     def estimate(