cholops 0.1.0__py3-none-any.whl

cholops/__init__.py

@@ -0,0 +1,28 @@
+"""Dense Cholesky modifications using ordinary JAX operations."""
+
+from .factors import Cholesky, PermutedCholesky
+from .tensors import (
+    chol_append,
+    chol_delete,
+    chol_delete_append,
+    chol_diag_update,
+    chol_rank1,
+    chol_rank1_positive,
+    chol_replace,
+    chol_replace_last,
+    cholesky,
+)
+
+__all__ = [
+    "Cholesky",
+    "PermutedCholesky",
+    "cholesky",
+    "chol_append",
+    "chol_delete",
+    "chol_diag_update",
+    "chol_rank1",
+    "chol_rank1_positive",
+    "chol_delete_append",
+    "chol_replace",
+    "chol_replace_last",
+]

cholops/factors/__init__.py

@@ -0,0 +1,6 @@
+"""Immutable Cholesky factor objects with matrix calculations and modifications."""
+
+from .cholesky import Cholesky
+from .permuted import PermutedCholesky
+
+__all__ = ["Cholesky", "PermutedCholesky"]

cholops/factors/_consumers.py

@@ -0,0 +1,28 @@
+"""Solves, log determinants and quadratic forms shared by the factor classes."""
+
+import jax.numpy as jnp
+from jax import Array
+from jax.scipy.linalg import solve_triangular
+
+
+def _rhs(L: Array, rhs: Array) -> Array:
+    rhs = jnp.asarray(rhs)
+    if rhs.ndim not in (1, 2) or rhs.shape[0] != L.shape[0]:
+        raise ValueError("rhs must have shape (n,) or (n,r)")
+    if rhs.dtype != L.dtype:
+        raise TypeError("rhs and the factor must have the same dtype")
+    return rhs
+
+
+def _solve(L: Array, rhs: Array) -> Array:
+    z = solve_triangular(L, rhs, lower=True)
+    return solve_triangular(L, z, lower=True, trans="T")
+
+
+def _logdet(L: Array) -> Array:
+    return 2 * jnp.sum(jnp.log(jnp.diag(L)))
+
+
+def _quadratic(L: Array, rhs: Array) -> Array:
+    z = solve_triangular(L, rhs, lower=True)
+    return jnp.sum(z * z, axis=0)

cholops/factors/cholesky.py

@@ -0,0 +1,133 @@
+"""Cholesky factor objects using the tensor operations."""
+
+from dataclasses import dataclass
+
+import jax
+import jax.numpy as jnp
+from jax import Array
+
+from ..tensors import (
+    chol_append,
+    chol_delete,
+    chol_delete_append,
+    chol_diag_update,
+    chol_rank1,
+    chol_rank1_positive,
+    chol_replace,
+    chol_replace_last,
+    cholesky,
+)
+from ..tensors._validation import _factor
+from ._consumers import _logdet, _quadratic, _rhs, _solve
+
+
+@jax.tree_util.register_dataclass
+@dataclass(frozen=True)
+class Cholesky:
+    """Canonical lower factor, stored as one dynamic PyTree leaf ``factor``.
+
+    Construct with from_matrix or from_factor. All inputs and outputs follow
+    the matrix's row and column order. Floating arrays share a real float32 or
+    float64 dtype. The represented matrix must remain positive definite.
+    Modifications return new objects, leaving the original unchanged.
+    Supports jit, vmap and first-order differentiation.
+    """
+
+    factor: Array
+
+    @classmethod
+    def from_matrix(cls, matrix: Array) -> "Cholesky":
+        """Factor real (n,n) input with the symmetry-preserving reverse rule.
+
+        The symmetric part must be positive definite. Batch with vmap.
+        """
+        matrix = jnp.asarray(matrix)
+        if matrix.ndim != 2:
+            raise ValueError("matrix must be unbatched; batch with jax.vmap")
+        return cls(cholesky(matrix))
+
+    @classmethod
+    def from_factor(cls, factor: Array) -> "Cholesky":
+        """Wrap a real (n,n) lower factor, ignoring its unused upper triangle.
+
+        Does not change the derivative of the function that produced factor.
+        Supply a valid factor with a positive diagonal. Empty factors are supported.
+        """
+        return cls(_factor(factor))
+
+    def solve(self, rhs: Array) -> Array:
+        """Solve the covariance system for a canonical (n,) or (n,r) RHS."""
+        return _solve(self.factor, _rhs(self.factor, rhs))
+
+    def logdet(self) -> Array:
+        """Scalar log determinant of the represented covariance."""
+        return _logdet(self.factor)
+
+    def quadratic(self, rhs: Array) -> Array:
+        """Inverse quadratic form: scalar for (n,), per column for (n,r)."""
+        return _quadratic(self.factor, _rhs(self.factor, rhs))
+
+    def append(self, k: Array, kappa: Array | float) -> "Cholesky":
+        """Append covariances k:(n,) and diagonal kappa; return size n+1.
+
+        The target covariance must be positive definite.
+        """
+        return type(self)(chol_append(self.factor, k, kappa))
+
+    def delete(self, index: Array | int) -> "Cholesky":
+        """Delete one variable, preserving survivor order; return size n-1.
+
+        index must be in [0,n). Check traced indices before calling this method,
+        since they are not checked at runtime.
+        """
+        return type(self)(chol_delete(self.factor, index))
+
+    def replace(self, index: Array | int, k: Array, kappa: Array | float) -> "Cholesky":
+        """Replace at index, preserving every other index and the (n,n) shape.
+
+        k:(n-1,) follows survivor order. kappa is the new diagonal.
+        The target must be positive definite and index must be in [0,n).
+        Traced indices are not checked at runtime.
+        """
+        return type(self)(chol_replace(self.factor, index, k, kappa))
+
+    def delete_append(
+        self, index: Array | int, k: Array, kappa: Array | float
+    ) -> "Cholesky":
+        """Delete index and append its replacement at the end, keeping size n.
+
+        k:(n-1,) follows survivor order. kappa is the new diagonal.
+        The returned factor is ordered as survivors followed by the replacement.
+        The target must be positive definite and index must be in [0,n).
+        Traced indices are not checked at runtime.
+        """
+        return type(self)(chol_delete_append(self.factor, index, k, kappa))
+
+    def replace_last(self, K: Array, C: Array) -> "Cholesky":
+        """Replace the last r variables with K:(n-r,r), C:(r,r), keeping size n.
+
+        Uses a Schur complement with native JAX Cholesky differentiation.
+        C must be symmetric and the target must be positive definite.
+        r is static, inferred from C; r=0 and r=n are supported.
+        """
+        return type(self)(chol_replace_last(self.factor, K, C))
+
+    def diag_update(self, index: Array | int, delta: Array | float) -> "Cholesky":
+        """Add signed delta to covariance diagonal index, preserving order.
+
+        The derivative with respect to delta is retained at delta=0.
+        The target must be positive definite.
+        """
+        return type(self)(chol_diag_update(self.factor, index, delta))
+
+    def rank1(self, u: Array, alpha: Array | float = 1.0) -> "Cholesky":
+        """Add alpha * outer(u,u), u:(n,), preserving order and size.
+
+        alpha may be positive, negative or zero. The weight derivative is
+        retained at zero. The target must be positive definite.
+        """
+        return type(self)(chol_rank1(self.factor, u, alpha))
+
+    def rank1_positive(self, u: Array) -> "Cholesky":
+        """Add outer(u,u), u:(n,), using the optimized positive-update algorithm."""
+        return type(self)(chol_rank1_positive(self.factor, u))

cholops/factors/permuted.py

@@ -0,0 +1,182 @@
+"""Cholesky factor objects that manage their internal variable ordering."""
+
+from dataclasses import dataclass
+
+import jax
+import jax.numpy as jnp
+from jax import Array
+
+from ..tensors import (
+    chol_append,
+    chol_delete,
+    chol_delete_append,
+    chol_diag_update,
+    chol_rank1,
+    chol_rank1_positive,
+    cholesky,
+)
+from ..tensors._indexing import _keep
+from ..tensors._validation import _factor, _index, _scalar, _vector
+from ._consumers import _logdet, _quadratic, _rhs, _solve
+
+
+@jax.tree_util.register_dataclass
+@dataclass(frozen=True)
+class PermutedCholesky:
+    """Physical lower factor and inverse permutations, all dynamic PyTree leaves.
+
+    order[physical] is logical; inverse_order[logical] is physical. Construct via
+    from_matrix or from_factor. All method inputs and outputs use logical order.
+    The stored factor is lower triangular in physical order. Floating arrays
+    share a real float32 or float64 dtype, and the represented matrix must remain
+    positive definite. Modifications return new objects.
+    Supports jit, vmap and first-order differentiation of continuous inputs.
+    """
+
+    physical_factor: Array
+    order: Array
+    inverse_order: Array
+
+    @classmethod
+    def from_matrix(cls, matrix: Array) -> "PermutedCholesky":
+        """Factor a real (n,n) matrix with the symmetry-preserving reverse rule.
+
+        The symmetric part must be positive definite. Inputs use float32 or
+        float64. Batch with vmap.
+        """
+        matrix = jnp.asarray(matrix)
+        if matrix.ndim != 2:
+            raise ValueError("matrix must be unbatched; batch with jax.vmap")
+        return cls.from_factor(cholesky(matrix))
+
+    @classmethod
+    def from_factor(cls, factor: Array) -> "PermutedCholesky":
+        """Wrap an existing (n,n) lower factor in identity logical order.
+
+        This does not change the derivative of the operation that produced
+        factor. Use from_matrix to apply cholops' Cholesky reverse rule.
+        """
+        factor = _factor(factor)
+        order = jnp.arange(factor.shape[0])
+        return cls(factor, order, order)
+
+    def append(self, k: Array, kappa: Array | float) -> "PermutedCholesky":
+        """Append a variable last in logical order, returning size n+1.
+
+        k:(n,) follows the existing logical order. kappa is the new diagonal.
+        """
+        L = self.physical_factor
+        n = L.shape[0]
+        k = _vector(k, n, L, "k")
+        result = chol_append(L, k[self.order], kappa)
+        last = jnp.asarray([n], dtype=self.order.dtype)
+        return type(self)(
+            result,
+            jnp.concatenate((self.order, last)),
+            jnp.concatenate((self.inverse_order, last)),
+        )
+
+    def delete(self, index: Array | int) -> "PermutedCholesky":
+        """Delete a logical variable, preserving survivor order and returning n-1.
+
+        Remaining logical indices are renumbered from zero.
+        """
+        L = self.physical_factor
+        n = L.shape[0]
+        index = _index(index, n)
+        physical = self.inverse_order[index]
+        result = chol_delete(L, physical)
+        survivors = self.order[_keep(n, physical)]
+        order = survivors - (survivors > index)
+        inverse = self.inverse_order[_keep(n, index)]
+        inverse = inverse - (inverse > physical)
+        return type(self)(result, order, inverse)
+
+    def replace(
+        self, index: Array | int, k: Array, kappa: Array | float
+    ) -> "PermutedCholesky":
+        """Return a new factor after replacing logical slot index.
+
+        k has shape (n-1,) in logical order with index omitted; scalar kappa
+        is the new covariance diagonal. Typed arrays share the factor dtype.
+        The target must be positive definite and index must be in [0,n).
+        Traced indices are not checked at runtime.
+        Logical order and shape are preserved.
+        """
+        L = self.physical_factor
+        n = L.shape[0]
+        index = _index(index, n)
+        k = _vector(k, n - 1, L, "k")
+        kappa = _scalar(kappa, L, "kappa")
+        physical = self.inverse_order[index]
+        keep = jnp.arange(n - 1) + (jnp.arange(n - 1) >= physical)
+        survivors = self.order[keep]
+        compact = survivors - (survivors > index)
+        result = chol_delete_append(L, physical, k[compact], kappa)
+        order = jnp.concatenate((survivors, index[None]))
+        inverse = self.inverse_order - (self.inverse_order > physical)
+        inverse = inverse.at[index].set(n - 1)
+        return type(self)(result, order, inverse)
+
+    def delete_append(
+        self, index: Array | int, k: Array, kappa: Array | float
+    ) -> "PermutedCholesky":
+        """Delete a logical variable and append its replacement last.
+
+        k:(n-1,) follows logical survivor order. kappa is the new diagonal.
+        Unlike replace, this changes logical order to survivors then replacement.
+        """
+        L = self.physical_factor
+        n = L.shape[0]
+        index = _index(index, n)
+        k = _vector(k, n - 1, L, "k")
+        physical = self.inverse_order[index]
+        survivors = self.order[_keep(n, physical)]
+        compact = survivors - (survivors > index)
+        result = chol_delete_append(L, physical, k[compact], kappa)
+        last = jnp.asarray([n - 1], dtype=self.order.dtype)
+        order = jnp.concatenate((compact, last))
+        inverse = self.inverse_order[_keep(n, index)]
+        inverse = jnp.concatenate((inverse - (inverse > physical), last))
+        return type(self)(result, order, inverse)
+
+    def diag_update(
+        self, index: Array | int, delta: Array | float
+    ) -> "PermutedCholesky":
+        """Add signed delta to one logical covariance diagonal entry."""
+        index = _index(index, self.physical_factor.shape[0])
+        result = chol_diag_update(
+            self.physical_factor, self.inverse_order[index], delta
+        )
+        return type(self)(result, self.order, self.inverse_order)
+
+    def rank1(self, u: Array, alpha: Array | float = 1.0) -> "PermutedCholesky":
+        """Add alpha * outer(u, u), with u:(n,) in logical order.
+
+        alpha may be positive, negative or zero. Logical ordering is unchanged.
+        """
+        L = self.physical_factor
+        u = _vector(u, L.shape[0], L, "u")
+        result = chol_rank1(L, u[self.order], alpha)
+        return type(self)(result, self.order, self.inverse_order)
+
+    def rank1_positive(self, u: Array) -> "PermutedCholesky":
+        """Add outer(u, u), with u:(n,) in logical order, using the fast update."""
+        L = self.physical_factor
+        u = _vector(u, L.shape[0], L, "u")
+        result = chol_rank1_positive(L, u[self.order])
+        return type(self)(result, self.order, self.inverse_order)
+
+    def solve(self, rhs: Array) -> Array:
+        """Solve the covariance system for canonical vector or matrix RHS."""
+        rhs = _rhs(self.physical_factor, rhs)
+        return _solve(self.physical_factor, rhs[self.order])[self.inverse_order]
+
+    def logdet(self) -> Array:
+        """Log determinant of the represented covariance."""
+        return _logdet(self.physical_factor)
+
+    def quadratic(self, rhs: Array) -> Array:
+        """Covariance inverse quadratic form, separately per matrix RHS column."""
+        rhs = _rhs(self.physical_factor, rhs)
+        return _quadratic(self.physical_factor, rhs[self.order])

cholops/tensors/__init__.py

@@ -0,0 +1,25 @@
+"""Differentiable Cholesky operations on dense lower-triangular arrays.
+
+Inputs use a common float32 or float64 dtype. Supply a valid lower factor and
+choose modifications that keep the represented matrix positive definite.
+Operations support JIT and vmap. Differentiation is supported only through
+first order, in both forward and reverse mode.
+"""
+
+from .append import chol_append
+from .deletion import chol_delete, chol_delete_append
+from .factorization import cholesky
+from .rank_one import chol_diag_update, chol_rank1, chol_rank1_positive
+from .replacement import chol_replace, chol_replace_last
+
+__all__ = [
+    "cholesky",
+    "chol_append",
+    "chol_delete",
+    "chol_diag_update",
+    "chol_rank1",
+    "chol_rank1_positive",
+    "chol_delete_append",
+    "chol_replace",
+    "chol_replace_last",
+]

cholops/tensors/_autodiff.py

@@ -0,0 +1,78 @@
+"""Custom JAX rules for first-order differentiation.
+
+Reuses each operation's reverse rule under JIT, batching and checkpointing.
+Only first-order differentiation is supported.
+"""
+
+import jax
+import jax.numpy as jnp
+from jax.experimental import hijax
+
+
+class _TangentMap(hijax.HiPrim):
+    def __init__(self, pullback, residuals, tangents, out_aval, pushforward=None):
+        self.in_avals = jax.tree.map(jax.typeof, (residuals, tangents))
+        self.out_aval = out_aval
+        self.params = dict(
+            pullback=pullback,
+            pushforward=pushforward,
+            in_avals=self.in_avals,
+            out_aval=out_aval,
+        )
+        super().__init__()
+
+    def expand(self, residuals, tangents):
+        if self.pushforward is not None:
+            return self.pushforward(residuals, tangents)
+        # Differentiating the linear pullback also handles masks introduced by
+        # batching runtime conditionals. Its primal at zero is discarded.
+        zero = jnp.zeros(self.out_aval.shape, self.out_aval.dtype)
+        _, pushforward = jax.vjp(lambda G: self.pullback(residuals, G), zero)
+        return pushforward(tangents)[0]
+
+    def transpose(self, cotangent, residuals, tangent_accums):
+        gradients = self.pullback(residuals, hijax.instantiate_zeros(cotangent))
+
+        def accumulate(accumulator, value):
+            if isinstance(accumulator, hijax.GradAccum):
+                accumulator.accum(value)
+
+        jax.tree.map(accumulate, tangent_accums, gradients)
+
+    def batch(self, axis_data, args, dims):
+        residuals, tangents = args
+        residual_dims, tangent_dims = dims
+
+        def pullback(residuals, G):
+            gradients = jax.vmap(
+                self.pullback, in_axes=(residual_dims, 0), axis_size=axis_data.size
+            )(residuals, G)
+            # A shared tangent is broadcast in the JVP, so its adjoint sums
+            # across the batch; mapped tangents retain their original axis.
+            return jax.tree.map(
+                lambda g, d: jnp.sum(g, axis=0) if d is None else jnp.moveaxis(g, 0, d),
+                gradients,
+                tangent_dims,
+                is_leaf=lambda x: x is None,
+            )
+
+        out_aval = self.out_aval.update(shape=(axis_data.size, *self.out_aval.shape))
+        pushforward = None
+        if self.pushforward is not None:
+
+            def pushforward(residuals, tangents):
+                return jax.vmap(
+                    self.pushforward,
+                    in_axes=(residual_dims, tangent_dims),
+                    axis_size=axis_data.size,
+                )(residuals, tangents)
+
+        operation = _TangentMap(pullback, residuals, tangents, out_aval, pushforward)
+        return operation(residuals, tangents), 0
+
+
+def tangent_map(pullback, residuals, tangents, output, pushforward=None):
+    """Apply the transpose of a linear pullback, retaining its explicit adjoint."""
+    return _TangentMap(pullback, residuals, tangents, jax.typeof(output), pushforward)(
+        residuals, tangents
+    )

cholops/tensors/_indexing.py

@@ -0,0 +1,51 @@
+"""Principal-submatrix indexing without quadratic Cartesian index tensors."""
+
+import jax
+import jax.numpy as jnp
+
+from ._autodiff import tangent_map
+
+
+def _keep(n, index):
+    positions = jnp.arange(n - 1)
+    return positions + (positions >= index)
+
+
+def _insert(matrix, index, row, column, diagonal):
+    """Insert one row/column with linear-sized gather indices.
+
+    Used both for factor assembly and the transpose of principal extraction.
+    The inserted diagonal is supplied separately from the off-diagonal vectors.
+    """
+    n = matrix.shape[0] + 1
+    if n == 1:
+        return jnp.reshape(diagonal, (1, 1))
+    positions = jnp.arange(n)
+    source = jnp.clip(positions - (positions > index), 0, n - 2)
+    result = jnp.take(jnp.take(matrix, source, axis=0), source, axis=1)
+    result = jnp.where((positions == index)[:, None], row[source][None, :], result)
+    result = jnp.where((positions == index)[None, :], column[source][:, None], result)
+    return jnp.where(
+        (positions == index)[:, None] & (positions == index)[None, :],
+        diagonal,
+        result,
+    )
+
+
+@jax.custom_jvp
+def _principal(matrix, index):
+    keep = _keep(matrix.shape[0], index)
+    return jnp.take(jnp.take(matrix, keep, axis=0), keep, axis=1)
+
+
+def _principal_pullback(index, cotangent):
+    zeros = jnp.zeros((cotangent.shape[0],), dtype=cotangent.dtype)
+    return _insert(cotangent, index, zeros, zeros, jnp.zeros((), cotangent.dtype))
+
+
+@_principal.defjvp
+def _principal_jvp(primals, tangents):
+    matrix, index = primals
+    dmatrix, _ = tangents
+    result = _principal(matrix, index)
+    return result, tangent_map(_principal_pullback, index, dmatrix, result)

cholops/tensors/_validation.py

@@ -0,0 +1,55 @@
+"""Shared shape, dtype and validity checks for tensor operations."""
+
+import jax.numpy as jnp
+from jax import Array
+
+
+def _factor(L: Array) -> Array:
+    L = jnp.asarray(L)
+    if L.ndim != 2 or L.shape[0] != L.shape[1]:
+        raise ValueError("L must be a square, unbatched matrix; batch with jax.vmap")
+    if L.dtype not in (jnp.float32, jnp.float64):
+        raise TypeError("L must have a real float32 or float64 dtype")
+    return jnp.tril(L)
+
+
+def _vector(x: Array, n: int, L: Array, name: str) -> Array:
+    x = jnp.asarray(x)
+    if x.shape != (n,):
+        raise ValueError(f"{name} must have shape ({n},), got {x.shape}")
+    if x.dtype != L.dtype:
+        raise TypeError(f"{name} and L must have the same dtype")
+    return x
+
+
+def _scalar(x: Array | float, L: Array, name: str) -> Array:
+    # Python scalars follow the factor precision; typed arrays are never cast.
+    x = jnp.asarray(x, dtype=L.dtype) if isinstance(x, (float, int)) else jnp.asarray(x)
+    if x.weak_type:
+        x = x.astype(L.dtype)
+    if x.shape != ():
+        raise ValueError(f"{name} must be a scalar")
+    if x.dtype != L.dtype:
+        raise TypeError(f"{name} and L must have the same dtype")
+    return x
+
+
+def _index(index: Array | int, n: int) -> Array:
+    if n == 0:
+        raise ValueError("cannot index an empty factor")
+    if isinstance(index, int) and not 0 <= index < n:
+        raise ValueError(f"index must be in [0, {n})")
+    index = jnp.asarray(index)
+    if index.shape != () or not jnp.issubdtype(index.dtype, jnp.integer):
+        raise TypeError("index must be an integer scalar")
+    return index
+
+
+def _valid_factor(L: Array) -> Array:
+    return jnp.all(jnp.isfinite(L)) & jnp.all(jnp.diag(L) > 0)
+
+
+def _finish(result: Array, valid: Array) -> Array:
+    valid = valid & _valid_factor(result)
+    # Keep the failure marker in arithmetic on the result, including its AD path.
+    return result * jnp.where(valid, 1, jnp.asarray(jnp.nan, dtype=result.dtype))

cholops/tensors/append.py

@@ -0,0 +1,30 @@
+"""Append a row and column through a bordering solve."""
+
+import jax.numpy as jnp
+from jax import Array
+from jax.scipy.linalg import solve_triangular
+
+from ._validation import _factor, _finish, _scalar, _valid_factor, _vector
+
+
+def _append(L: Array, k: Array, kappa: Array) -> Array:
+    n = L.shape[0]
+    v = solve_triangular(L, k, lower=True) if n else k
+    diagonal = jnp.sqrt(kappa - jnp.vdot(v, v))
+    result = jnp.zeros((n + 1, n + 1), dtype=L.dtype)
+    return result.at[:n, :n].set(L).at[n, :n].set(v).at[n, n].set(diagonal)
+
+
+def chol_append(L: Array, k: Array, kappa: Array | float) -> Array:
+    """Factor [[L @ L.T, k], [k.T, kappa]], returning shape (n+1, n+1).
+
+    k has shape (n,) in the current physical ordering; kappa is scalar.
+    The Schur complement kappa - ||solve(L, k)||^2 must be strictly positive.
+    Empty (0, 0) factors are supported. Work is O(n^2).
+    Supports jit, vmap, JVP and reverse-mode AD. Numerical failure gives NaNs.
+    """
+    L = _factor(L)
+    k = _vector(k, L.shape[0], L, "k")
+    kappa = _scalar(kappa, L, "kappa")
+    valid = _valid_factor(L) & jnp.all(jnp.isfinite(k)) & jnp.isfinite(kappa)
+    return _finish(_append(L, k, kappa), valid)