gpjax 0.11.2__py3-none-any.whl → 0.12.2__py3-none-any.whl

gpjax/linalg/operators.py

@@ -0,0 +1,411 @@
+"""Linear operator abstractions for GPJax."""
+
+from abc import (
+    ABC,
+    abstractmethod,
+)
+from typing import (
+    Any,
+    List,
+    Tuple,
+    Union,
+)
+
+from jax import Array
+import jax.numpy as jnp
+import jax.tree_util as jtu
+from jaxtyping import Float
+
+
+class LinearOperator(ABC):
+    """Abstract base class for linear operators."""
+
+    def __init__(self):
+        super().__init__()
+
+    @property
+    @abstractmethod
+    def shape(self) -> Tuple[int, int]:
+        """Return the shape of the operator."""
+
+    @property
+    @abstractmethod
+    def dtype(self) -> jnp.dtype:
+        """Return the data type of the operator."""
+
+    @abstractmethod
+    def to_dense(self) -> Float[Array, "M N"]:
+        """Convert the operator to a dense JAX array."""
+
+    @property
+    def T(self) -> "LinearOperator":
+        """Return the transpose of the operator."""
+        # Default implementation: convert to dense and transpose
+        return Dense(self.to_dense().T)
+
+    def __matmul__(self, other):
+        """Matrix multiplication with another array or operator."""
+        if hasattr(other, "to_dense"):
+            # Other is a LinearOperator
+            return Dense(self.to_dense() @ other.to_dense())
+        else:
+            # Other is a JAX array
+            return self.to_dense() @ other
+
+    def __rmatmul__(self, other):
+        """Right matrix multiplication (other @ self)."""
+        if hasattr(other, "to_dense"):
+            # Other is a LinearOperator
+            return Dense(other.to_dense() @ self.to_dense())
+        else:
+            # Other is a JAX array
+            return other @ self.to_dense()
+
+    def __add__(self, other):
+        """Addition with another array or operator."""
+        if hasattr(other, "to_dense"):
+            # Other is a LinearOperator
+            return Dense(self.to_dense() + other.to_dense())
+        else:
+            # Other is a JAX array
+            return Dense(self.to_dense() + other)
+
+    def __radd__(self, other):
+        """Right addition (other + self)."""
+        if hasattr(other, "to_dense"):
+            # Other is a LinearOperator
+            return Dense(other.to_dense() + self.to_dense())
+        else:
+            # Other is a JAX array
+            return Dense(other + self.to_dense())
+
+    def __sub__(self, other):
+        """Subtraction with another array or operator."""
+        if hasattr(other, "to_dense"):
+            # Other is a LinearOperator
+            return Dense(self.to_dense() - other.to_dense())
+        else:
+            # Other is a JAX array
+            return Dense(self.to_dense() - other)
+
+    def __rsub__(self, other):
+        """Right subtraction (other - self)."""
+        if hasattr(other, "to_dense"):
+            # Other is a LinearOperator
+            return Dense(other.to_dense() - self.to_dense())
+        else:
+            # Other is a JAX array
+            return Dense(other - self.to_dense())
+
+    def __mul__(self, other):
+        """Scalar multiplication (self * scalar)."""
+        if jnp.isscalar(other):
+            return Dense(self.to_dense() * other)
+        else:
+            # Element-wise multiplication with array
+            return Dense(self.to_dense() * other)
+
+    def __rmul__(self, other):
+        """Right scalar multiplication (scalar * self)."""
+        if jnp.isscalar(other):
+            return Dense(other * self.to_dense())
+        else:
+            # Element-wise multiplication with array
+            return Dense(other * self.to_dense())
+
+
+class Dense(LinearOperator):
+    """Dense linear operator wrapping a JAX array."""
+
+    def __init__(self, array: Float[Array, "M N"]):
+        super().__init__()
+        self.array = array
+
+    @property
+    def shape(self) -> Tuple[int, int]:
+        return self.array.shape
+
+    @property
+    def dtype(self) -> jnp.dtype:
+        return self.array.dtype
+
+    def to_dense(self) -> Float[Array, "M N"]:
+        return self.array
+
+    @property
+    def T(self) -> "Dense":
+        return Dense(self.array.T)
+
+
+class Diagonal(LinearOperator):
+    """Diagonal linear operator."""
+
+    def __init__(self, diagonal: Float[Array, " N"]):
+        super().__init__()
+        self.diagonal = diagonal
+
+    @property
+    def shape(self) -> Tuple[int, int]:
+        n = self.diagonal.shape[0]
+        return (n, n)
+
+    @property
+    def dtype(self) -> jnp.dtype:
+        return self.diagonal.dtype
+
+    def to_dense(self) -> Float[Array, "N N"]:
+        return jnp.diag(self.diagonal)
+
+    @property
+    def T(self) -> "Diagonal":
+        return Diagonal(self.diagonal)
+
+
+class Identity(LinearOperator):
+    """Identity linear operator."""
+
+    def __init__(self, shape: Union[int, Tuple[int, int]], dtype=jnp.float64):
+        super().__init__()
+        if isinstance(shape, int):
+            self._shape = (shape, shape)
+        else:
+            if shape[0] != shape[1]:
+                raise ValueError(f"Identity matrix must be square, got shape {shape}")
+            self._shape = shape
+        self._dtype = dtype
+
+    @property
+    def shape(self) -> Tuple[int, int]:
+        return self._shape
+
+    @property
+    def dtype(self) -> Any:
+        return self._dtype
+
+    def to_dense(self) -> Float[Array, "N N"]:
+        n = self._shape[0]
+        return jnp.eye(n, dtype=self._dtype)
+
+    @property
+    def T(self) -> "Identity":
+        return Identity(self._shape, dtype=self._dtype)
+
+
+class Triangular(LinearOperator):
+    """Triangular linear operator."""
+
+    def __init__(self, array: Float[Array, "N N"], lower: bool = True):
+        super().__init__()
+        self.array = array
+        self.lower = lower
+
+    @property
+    def shape(self) -> Tuple[int, int]:
+        return self.array.shape
+
+    @property
+    def dtype(self) -> Any:
+        return self.array.dtype
+
+    def to_dense(self) -> Float[Array, "N N"]:
+        if self.lower:
+            return jnp.tril(self.array)
+        else:
+            return jnp.triu(self.array)
+
+    @property
+    def T(self) -> "Triangular":
+        return Triangular(self.array.T, lower=not self.lower)
+
+
+class BlockDiag(LinearOperator):
+    """Block diagonal linear operator."""
+
+    def __init__(
+        self, operators: List[LinearOperator], multiplicities: List[int] = None
+    ):
+        super().__init__()
+        self.operators = operators
+
+        # Handle multiplicities - how many times each block is repeated
+        if multiplicities is None:
+            self.multiplicities = [1] * len(operators)
+        else:
+            if len(multiplicities) != len(operators):
+                raise ValueError(
+                    f"Length of multiplicities ({len(multiplicities)}) must match operators ({len(operators)})"
+                )
+            self.multiplicities = multiplicities
+
+        # Calculate total shape with multiplicities
+        rows = sum(
+            op.shape[0] * mult
+            for op, mult in zip(operators, self.multiplicities, strict=False)
+        )
+        cols = sum(
+            op.shape[1] * mult
+            for op, mult in zip(operators, self.multiplicities, strict=False)
+        )
+        self._shape = (rows, cols)
+
+        # Use dtype of first operator (assuming all same dtype)
+        if operators:
+            self._dtype = operators[0].dtype
+        else:
+            self._dtype = jnp.float64
+
+    @property
+    def shape(self) -> Tuple[int, int]:
+        return self._shape
+
+    @property
+    def dtype(self) -> Any:
+        return self._dtype
+
+    def to_dense(self) -> Float[Array, "M N"]:
+        if not self.operators:
+            return jnp.zeros(self._shape, dtype=self._dtype)
+
+        # Convert each operator to dense and create block diagonal with multiplicities
+        expanded_blocks = []
+        for op, mult in zip(self.operators, self.multiplicities, strict=False):
+            op_dense = op.to_dense()
+            for _ in range(mult):
+                expanded_blocks.append(op_dense)
+
+        # Create the full block diagonal matrix
+        n_blocks = len(expanded_blocks)
+        if n_blocks == 0:
+            return jnp.zeros(self._shape, dtype=self._dtype)
+
+        # Build the block diagonal matrix
+        rows = []
+        for i in range(n_blocks):
+            row = []
+            for j in range(n_blocks):
+                if i == j:
+                    row.append(expanded_blocks[i])
+                else:
+                    row.append(
+                        jnp.zeros(
+                            (expanded_blocks[i].shape[0], expanded_blocks[j].shape[1]),
+                            dtype=self._dtype,
+                        )
+                    )
+            rows.append(row)
+        return jnp.block(rows)
+
+    @property
+    def T(self) -> "BlockDiag":
+        transposed_ops = [op.T for op in self.operators]
+        return BlockDiag(transposed_ops, multiplicities=self.multiplicities)
+
+
+class Kronecker(LinearOperator):
+    """Kronecker product linear operator."""
+
+    def __init__(self, operators: List[LinearOperator]):
+        super().__init__()
+        if len(operators) < 2:
+            raise ValueError("Kronecker product requires at least 2 operators")
+        self.operators = operators
+
+        # Calculate shape as product of individual shapes
+        rows = 1
+        cols = 1
+        for op in operators:
+            rows *= op.shape[0]
+            cols *= op.shape[1]
+        self._shape = (rows, cols)
+
+        # Use dtype of first operator
+        self._dtype = operators[0].dtype
+
+    @property
+    def shape(self) -> Tuple[int, int]:
+        return self._shape
+
+    @property
+    def dtype(self) -> Any:
+        return self._dtype
+
+    def to_dense(self) -> Float[Array, "M N"]:
+        # Convert to dense and compute Kronecker product
+        result = self.operators[0].to_dense()
+        for op in self.operators[1:]:
+            result = jnp.kron(result, op.to_dense())
+        return result
+
+    @property
+    def T(self) -> "Kronecker":
+        transposed_ops = [op.T for op in self.operators]
+        return Kronecker(transposed_ops)
+
+
+def _dense_tree_flatten(dense):
+    return (dense.array,), None
+
+
+def _dense_tree_unflatten(aux_data, children):
+    return Dense(children[0])
+
+
+jtu.register_pytree_node(Dense, _dense_tree_flatten, _dense_tree_unflatten)
+
+
+def _diagonal_tree_flatten(diagonal):
+    return (diagonal.diagonal,), None
+
+
+def _diagonal_tree_unflatten(aux_data, children):
+    return Diagonal(children[0])
+
+
+jtu.register_pytree_node(Diagonal, _diagonal_tree_flatten, _diagonal_tree_unflatten)
+
+
+def _identity_tree_flatten(identity):
+    return (), (identity._shape, identity._dtype)
+
+
+def _identity_tree_unflatten(aux_data, children):
+    shape, dtype = aux_data
+    return Identity(shape, dtype)
+
+
+jtu.register_pytree_node(Identity, _identity_tree_flatten, _identity_tree_unflatten)
+
+
+def _triangular_tree_flatten(triangular):
+    return (triangular.array,), triangular.lower
+
+
+def _triangular_tree_unflatten(aux_data, children):
+    return Triangular(children[0], aux_data)
+
+
+jtu.register_pytree_node(
+    Triangular, _triangular_tree_flatten, _triangular_tree_unflatten
+)
+
+
+def _blockdiag_tree_flatten(blockdiag):
+    return tuple(blockdiag.operators), blockdiag.multiplicities
+
+
+def _blockdiag_tree_unflatten(aux_data, children):
+    return BlockDiag(list(children), aux_data)
+
+
+jtu.register_pytree_node(BlockDiag, _blockdiag_tree_flatten, _blockdiag_tree_unflatten)
+
+
+def _kronecker_tree_flatten(kronecker):
+    return tuple(kronecker.operators), None
+
+
+def _kronecker_tree_unflatten(aux_data, children):
+    return Kronecker(list(children))
+
+
+jtu.register_pytree_node(Kronecker, _kronecker_tree_flatten, _kronecker_tree_unflatten)

gpjax/linalg/utils.py

@@ -0,0 +1,65 @@
+"""Utility functions for the linear algebra module."""
+
+import jax.numpy as jnp
+from jaxtyping import Array
+
+from gpjax.linalg.operators import LinearOperator
+
+
+class PSDAnnotation:
+    """Marker class for PSD (Positive Semi-Definite) annotations."""
+
+    def __call__(self, A: LinearOperator) -> LinearOperator:
+        """Make PSD annotation callable."""
+        return psd(A)
+
+
+# Create the PSD marker similar to cola.PSD
+PSD = PSDAnnotation()
+
+
+def psd(A: LinearOperator) -> LinearOperator:
+    """Mark a linear operator as positive semi-definite.
+
+    This function acts as a marker/wrapper for positive semi-definite matrices.
+
+    Args:
+        A: A LinearOperator that is assumed to be positive semi-definite.
+
+    Returns:
+        The same LinearOperator, marked as PSD.
+    """
+    # Add annotations attribute if it doesn't exist
+    if not hasattr(A, "annotations"):
+        A.annotations = set()
+    A.annotations.add(PSD)
+    return A
+
+
+def add_jitter(matrix: Array, jitter: float | Array = 1e-6) -> Array:
+    """Add jitter to the diagonal of a matrix for numerical stability.
+
+    This function adds a small positive value (jitter) to the diagonal elements
+    of a square matrix to improve numerical stability, particularly for
+    Cholesky decompositions and matrix inversions.
+
+    Args:
+        matrix: A square matrix to which jitter will be added.
+        jitter: The jitter value to add to the diagonal. Defaults to 1e-6.
+
+    Returns:
+        The matrix with jitter added to its diagonal.
+
+    Examples:
+        >>> import jax.numpy as jnp
+        >>> from gpjax.linalg.utils import add_jitter
+        >>> matrix = jnp.array([[1.0, 0.5], [0.5, 1.0]])
+        >>> jittered_matrix = add_jitter(matrix, jitter=0.01)
+    """
+    if matrix.ndim != 2:
+        raise ValueError(f"Expected 2D matrix, got {matrix.ndim}D array")
+
+    if matrix.shape[0] != matrix.shape[1]:
+        raise ValueError(f"Expected square matrix, got shape {matrix.shape}")
+
+    return matrix + jnp.eye(matrix.shape[0]) * jitter

gpjax/mean_functions.py

@@ -27,8 +27,6 @@ from jaxtyping import (
 
 from gpjax.parameters import (
     Parameter,
-    Real,
-    Static,
 )
 from gpjax.typing import (
     Array,
@@ -132,12 +130,12 @@ class Constant(AbstractMeanFunction):
 
     def __init__(
         self,
-        constant: tp.Union[ScalarFloat, Float[Array, " O"], Parameter, Static] = 0.0,
+        constant: tp.Union[ScalarFloat, Float[Array, " O"], Parameter] = 0.0,
     ):
-        if isinstance(constant, Parameter) or isinstance(constant, Static):
+        if isinstance(constant, Parameter):
             self.constant = constant
         else:
-            self.constant = Real(jnp.array(constant))
+            self.constant = jnp.array(constant)
 
     def __call__(self, x: Num[Array, "N D"]) -> Float[Array, "N O"]:
         r"""Evaluate the mean function at the given points.
@@ -148,7 +146,10 @@ class Constant(AbstractMeanFunction):
         Returns:
             Float[Array, "1"]: The evaluated mean function.
         """
-        return jnp.ones((x.shape[0], 1)) * self.constant.value
+        if isinstance(self.constant, Parameter):
+            return jnp.ones((x.shape[0], 1)) * self.constant.value
+        else:
+            return jnp.ones((x.shape[0], 1)) * self.constant
 
 
 class Zero(Constant):
@@ -160,7 +161,7 @@ class Zero(Constant):
     """
 
     def __init__(self):
-        super().__init__(constant=Static(jnp.array(0.0)))
+        super().__init__(constant=0.0)
 
 
 class CombinationMeanFunction(AbstractMeanFunction):

gpjax/objectives.py

@@ -1,13 +1,5 @@
 from typing import TypeVar
 
-from cola.annotations import PSD
-from cola.linalg.decompositions.decompositions import Cholesky
-from cola.linalg.inverse.inv import (
-    inv,
-    solve,
-)
-from cola.linalg.trace.diag_trace import diag
-from cola.ops.operators import I_like
 from flax import nnx
 from jax import vmap
 import jax.numpy as jnp
@@ -22,7 +14,13 @@ from gpjax.gps import (
     ConjugatePosterior,
     NonConjugatePosterior,
 )
-from gpjax.lower_cholesky import lower_cholesky
+from gpjax.linalg import (
+    Dense,
+    lower_cholesky,
+    psd,
+    solve,
+)
+from gpjax.linalg.utils import add_jitter
 from gpjax.typing import (
     Array,
     ScalarFloat,
@@ -100,9 +98,9 @@ def conjugate_mll(posterior: ConjugatePosterior, data: Dataset) -> ScalarFloat:
 
     # Σ = (Kxx + Io²) = LLᵀ
     Kxx = posterior.prior.kernel.gram(x)
-    Kxx += I_like(Kxx) * posterior.prior.jitter
-    Sigma = Kxx + I_like(Kxx) * obs_noise
-    Sigma = PSD(Sigma)
+    Kxx_dense = add_jitter(Kxx.to_dense(), posterior.prior.jitter)
+    Sigma_dense = Kxx_dense + jnp.eye(Kxx.shape[0]) * obs_noise
+    Sigma = psd(Dense(Sigma_dense))
 
     # p(y | x, θ), where θ are the model hyperparameters:
     mll = GaussianDistribution(jnp.atleast_1d(mx.squeeze()), Sigma)
@@ -164,11 +162,14 @@ def conjugate_loocv(posterior: ConjugatePosterior, data: Dataset) -> ScalarFloat
 
     # Σ = (Kxx + Io²)
     Kxx = posterior.prior.kernel.gram(x)
-    Sigma = Kxx + I_like(Kxx) * (obs_var + posterior.prior.jitter)
-    Sigma = PSD(Sigma)  # [N, N]
+    Sigma_dense = Kxx.to_dense() + jnp.eye(Kxx.shape[0]) * (
+        obs_var + posterior.prior.jitter
+    )
+    Sigma = psd(Dense(Sigma_dense))  # [N, N]
 
-    Sigma_inv_y = solve(Sigma, y - mx, Cholesky())  # [N, 1]
-    Sigma_inv_diag = diag(inv(Sigma, Cholesky()))[:, None]  # [N, 1]
+    Sigma_inv_y = solve(Sigma, y - mx)  # [N, 1]
+    Sigma_inv = jnp.linalg.inv(Sigma.to_dense())
+    Sigma_inv_diag = jnp.diag(Sigma_inv)[:, None]  # [N, 1]
 
     loocv_means = mx + (y - mx) - Sigma_inv_y / Sigma_inv_diag
     loocv_stds = jnp.sqrt(1.0 / Sigma_inv_diag)
@@ -213,8 +214,8 @@ def log_posterior_density(
 
     # Gram matrix
     Kxx = posterior.prior.kernel.gram(x)
-    Kxx += I_like(Kxx) * posterior.prior.jitter
-    Kxx = PSD(Kxx)
+    Kxx_dense = add_jitter(Kxx.to_dense(), posterior.prior.jitter)
+    Kxx = psd(Dense(Kxx_dense))
     Lx = lower_cholesky(Kxx)
 
     # Compute the prior mean function
@@ -349,8 +350,8 @@ def collapsed_elbo(variational_family: VF, data: Dataset) -> ScalarFloat:
     noise = variational_family.posterior.likelihood.obs_stddev.value**2
     z = variational_family.inducing_inputs.value
     Kzz = kernel.gram(z)
-    Kzz += I_like(Kzz) * variational_family.jitter
-    Kzz = PSD(Kzz)
+    Kzz_dense = add_jitter(Kzz.to_dense(), variational_family.jitter)
+    Kzz = psd(Dense(Kzz_dense))
     Kzx = kernel.cross_covariance(z, x)
     Kxx_diag = vmap(kernel, in_axes=(0, 0))(x, x)
     μx = mean_function(x)
@@ -383,7 +384,7 @@ def collapsed_elbo(variational_family: VF, data: Dataset) -> ScalarFloat:
     #
     #   with A and B defined as above.
 
-    A = solve(Lz, Kzx, Cholesky()) / jnp.sqrt(noise)
+    A = solve(Lz, Kzx) / jnp.sqrt(noise)
 
     # AAᵀ
     AAT = jnp.matmul(A, A.T)

gpjax/parameters.py

@@ -21,23 +21,20 @@ def transform(
     r"""Transforms parameters using a bijector.
 
     Example:
-    ```pycon
         >>> from gpjax.parameters import PositiveReal, transform
         >>> import jax.numpy as jnp
         >>> import numpyro.distributions.transforms as npt
         >>> from flax import nnx
         >>> params = nnx.State(
-        >>>     {
-        >>>         "a": PositiveReal(jnp.array([1.0])),
-        >>>         "b": PositiveReal(jnp.array([2.0])),
-        >>>     }
-        >>> )
+        ...     {
+        ...         "a": PositiveReal(jnp.array([1.0])),
+        ...         "b": PositiveReal(jnp.array([2.0])),
+        ...     }
+        ... )
         >>> params_bijection = {'positive': npt.SoftplusTransform()}
         >>> transformed_params = transform(params, params_bijection)
         >>> print(transformed_params["a"].value)
-         [1.3132617]
-    ```
-
+        [1.3132617]
 
     Args:
         params: A nnx.State object containing parameters to be transformed.
@@ -49,7 +46,7 @@ def transform(
     """
 
     def _inner(param):
-        bijector = params_bijection.get(param._tag, npt.IdentityTransform())
+        bijector = params_bijection.get(param.tag, npt.IdentityTransform())
         if inverse:
             transformed_value = bijector.inv(param.value)
         else:
@@ -60,10 +57,11 @@ def transform(
 
     gp_params, *other_params = params.split(Parameter, ...)
 
+    # Transform each parameter in the state
     transformed_gp_params: nnx.State = jtu.tree_map(
-        lambda x: _inner(x),
+        lambda x: _inner(x) if isinstance(x, Parameter) else x,
         gp_params,
-        is_leaf=lambda x: isinstance(x, nnx.VariableState),
+        is_leaf=lambda x: isinstance(x, Parameter),
     )
     return nnx.State.merge(transformed_gp_params, *other_params)
 
@@ -79,7 +77,7 @@ class Parameter(nnx.Variable[T]):
         _check_is_arraylike(value)
 
         super().__init__(value=jnp.asarray(value), **kwargs)
-        self._tag = tag
+        self.tag = tag
 
 
 class NonNegativeReal(Parameter[T]):
@@ -124,16 +122,6 @@ class SigmoidBounded(Parameter[T]):
             )
 
 
-class Static(nnx.Variable[T]):
-    """Static parameter that is not trainable."""
-
-    def __init__(self, value: T, tag: ParameterTag = "static", **kwargs):
-        _check_is_arraylike(value)
-
-        super().__init__(value=jnp.asarray(value), tag=tag, **kwargs)
-        self._tag = tag
-
-
 class LowerTriangular(Parameter[T]):
     """Parameter that is a lower triangular matrix."""