lrux 0.1.0__py3-none-any.whl

lrux/__init__.py

@@ -0,0 +1,3 @@
+from .det_lru import det_lru, DetCarrier, init_det_carrier, det_lru_delayed
+from .pfaffian import skew_eye, pf, slogpf
+from .pf_lru import pf_lru, PfCarrier, init_pf_carrier, pf_lru_delayed

lrux/det_lru.py

@@ -0,0 +1,563 @@
+from typing import Optional, Tuple, Union, Sequence, NamedTuple
+from jax import Array
+from jax.typing import ArrayLike
+import jax
+import jax.numpy as jnp
+from jax._src.numpy import reductions
+
+
+def _check_mat(mat: Array) -> None:
+    if mat.ndim != 2 or mat.shape[0] != mat.shape[1]:
+        raise ValueError(f"Expect input matrix shape (n, n), got {mat.shape}.")
+
+
+def _standardize_uv(
+    u: Union[ArrayLike, Tuple[Array, ArrayLike]], n: int, dtype: jnp.dtype
+) -> Tuple[Array, Array]:
+    if isinstance(u, ArrayLike):
+        u = jnp.asarray(u)
+        if jnp.issubdtype(u.dtype, jnp.integer):
+            u = (jnp.empty((n, 0), dtype), u.flatten())
+        else:
+            u = (u.reshape(n, -1), jnp.array([], dtype=jnp.int32))
+    elif isinstance(u, Sequence):
+        u = (jnp.asarray(u[0]).reshape(n, -1), jnp.asarray(u[1]).flatten())
+    else:
+        raise ValueError(f"Got unsupported u or v data type {type(u)}.")
+    return u
+
+
+def _check_uv(u: Tuple[Array, Array], v: Tuple[Array, Array]) -> None:
+    rank_u = u[0].shape[1] + u[1].size
+    rank_v = v[0].shape[1] + v[1].size
+    if rank_u != rank_v:
+        raise ValueError(
+            f"The input u and v should have matched rank, got {rank_u} and {rank_v}."
+        )
+
+
+def _get_R(Ainv: Array, u: Tuple[Array, Array], v: Tuple[Array, Array]) -> Array:
+    xu_Ainv_xv = jnp.einsum("nk,nm,ml->kl", u[0], Ainv, v[0])
+    eu_Ainv_xv = Ainv[u[1]] @ v[0]
+    xu_Ainv_ev = u[0].T @ Ainv[:, v[1]]
+    eu_Ainv_ev = Ainv[u[1], :][:, v[1]]
+    uT_Ainv_v = jnp.block([[xu_Ainv_ev, xu_Ainv_xv], [eu_Ainv_ev, eu_Ainv_xv]])
+    return uT_Ainv_v.at[jnp.diag_indices_from(uT_Ainv_v)].add(1)
+
+
+def _det_and_lufac(R: Array) -> Tuple[Array, Tuple[Array, Array]]:
+    lu, pivot = jax.scipy.linalg.lu_factor(R)
+    iota = jnp.arange(pivot.size, dtype=pivot.dtype)
+    parity = reductions.count_nonzero(pivot != iota, axis=-1)
+    sign = jnp.array(-2 * (parity % 2) + 1, dtype=lu.dtype)
+    det = sign * jnp.prod(jnp.diag(lu))
+    return det, (lu, pivot)
+
+
+def _update_Ainv(
+    Ainv: Array,
+    u: Tuple[Array, Array],
+    v: Tuple[Array, Array],
+    lu_and_piv: Tuple[Array, Array],
+) -> Array:
+    uT_Ainv = jnp.concatenate((u[0].T @ Ainv, Ainv[u[1], :]), axis=0)
+    Rinv_uT_Ainv = jax.scipy.linalg.lu_solve(lu_and_piv, uT_Ainv)
+    Ainv_v = jnp.concatenate((Ainv[:, v[1]], Ainv @ v[0]), axis=1)
+    return Ainv - Ainv_v @ Rinv_uT_Ainv
+
+
+def det_lru(
+    Ainv: Array,
+    u: Union[ArrayLike, Tuple[Array, ArrayLike]],
+    v: Union[ArrayLike, Tuple[Array, ArrayLike]],
+    return_update: bool = False,
+) -> Union[Array, Tuple[Array, Array]]:
+    r"""
+    Low-rank update of determinant :math:`\det(A_1) = \det(A_0 + vu^T)`
+
+    :param Ainv:
+        Inverse of the original matrix :math:`A_0^{-1}`, shape (n, n)
+
+    :param u:
+        Low-rank update vector(s) :math:`u`. There are several acceptable inputs
+        of ``u`` as listed below.
+
+        An array with shape (n,) or (n, k):
+            Direct expression of full low-rank vector(s).
+
+        An integer or array of integers with size k:
+            One-hot vectors ``u_full = jnp.zeros((n, k)).at[u, jnp.arange(k)].set(1)``.
+            For example, when you need a full matrix
+
+            .. code-block:: python
+
+                u = jnp.array([
+                    [0, 0],
+                    [1, 0], 
+                    [0, 0], 
+                    [0, 1],
+                ])
+            
+            you can alternatively specify 
+
+            .. code-block:: python
+            
+                u = jnp.array([1, 3])
+
+        A tuple of two arrays, with respective shapes (n, k0) and (k1,):
+            A concatenation of the previous two. For example, when you need a full matrix
+
+            .. code-block:: python
+
+                u = jnp.array([
+                    [u00, u01, 0, 0],
+                    [u10, u11, 1, 0], 
+                    [u20, u21, 0, 0], 
+                    [u30, u31, 0, 1],
+                ])
+
+            you can alternatively specify
+
+            .. code-block:: python
+            
+                x = jnp.array([
+                    [u00, u01],
+                    [u10, u11], 
+                    [u20, u21], 
+                    [u30, u31],
+                ])
+                e = jnp.array([1, 3])
+                u = (x, e)
+
+        The matrix product of one-hot vectors is internally performed by matrix slicing
+        for better performance, so an input of indices is preferred.
+
+    :param v:
+        Low-rank update vector(s) :math:`v`. The acceptable inputs are similar to ``u``.
+        When the input is a tuple of two arrays ``v = (x, e)``, for convenience
+        the concatenation order is reversely given by ``jnp.concatenate((ve, vx), axis=1)``.
+        Therefore, when you need
+
+        .. code-block:: python
+
+            v = jnp.array([
+                [0, 0, v00, v01],
+                [1, 0, v10, v11], 
+                [0, 0, v20, u21], 
+                [0, 1, v30, v31],
+            ])
+
+        you can alternatively specify
+
+        .. code-block:: python
+        
+            x = jnp.array([
+                [v00, v01],
+                [v10, v11], 
+                [v20, v21], 
+                [v30, v31],
+            ])
+            e = jnp.array([1, 3])
+            u = (x, e)
+
+    :param return_update:
+        Whether the new matrix inverse :math:`A_1^{-1}` should be returned,
+        defaul to False.
+
+    :return:
+        ratio:
+            The ratio between two determinants
+
+            .. math::
+
+                r = \frac{\det(A_1)}{\det(A_0)} = \det(R)
+
+            where
+
+            .. math::
+
+                R = I + u^T A_0^{-1} v
+
+        new_Ainv:
+            The new matrix inverse
+
+            .. math::
+
+                A_1^{-1} = (A_0 + vu^T)^{-1} = A_0^{-1} - A_0^{-1} v R^{-1} u^T A_0^{-1}
+
+            Only returned when ``return_update`` is True.
+
+    .. tip::
+
+        This function is compatible with ``jax.jit`` and ``jax.vmap``, while
+        ``return_update`` is a static argument which shouldn't be jitted or vmapped.
+
+        Furthermore, we recommend setting ``donate_argnums=0`` in ``jax.jit`` to reuse 
+        the memory of ``Ainv`` if it's no longer needed. This helps to greatly reduce 
+        the time and memory cost. For instance,
+
+        .. code-block:: python
+
+            lru_vmap = jax.vmap(det_lru, in_axes=(0, 0, 0, None))
+            lru_jit = jax.jit(lru_vmap, static_argnums=3, donate_argnums=0)
+
+    .. note::
+
+        Here are examples of how to define ``u`` and ``v`` before calling ``det_lru(Ainv, u, v)``.
+        Keep in mind that the low-rank update we need takes the form
+
+        .. math::
+
+            A_1 - A_0 = vu^T
+
+        **Rank-1 row update**
+
+        .. math::
+
+            A_1 - A_0 = \begin{pmatrix}
+                0 & 0 & 0 & 0 \\ 
+                u_0 & u_1 & u_2 & u_3 \\
+                0 & 0 & 0 & 0 \\ 
+                0 & 0 & 0 & 0 \\ 
+            \end{pmatrix}
+            = \begin{pmatrix}
+                0 \\ 1 \\ 0 \\ 0
+            \end{pmatrix}
+            (u_0, u_1, u_2, u_3)
+            
+        .. code-block:: python
+        
+            u = jnp.array([u0, u1, u2, u3])
+            v = 1
+
+        **Rank-1 column update**
+
+        .. math::
+
+            A_1 - A_0 = \begin{pmatrix}
+                0 & 0 & v_0 & 0 \\ 
+                0 & 0 & v_1 & 0 \\
+                0 & 0 & v_2 & 0 \\ 
+                0 & 0 & v_3 & 0 \\ 
+            \end{pmatrix}
+            = \begin{pmatrix}
+                v_0 \\ v_1 \\ v_2 \\ v_3
+            \end{pmatrix}
+            (0, 0, 1, 0)
+            
+        .. code-block:: python
+        
+            u = 2
+            v = jnp.array([v0, v1, v2, v3])
+
+        **Rank-2 row update**
+
+        .. math::
+
+            A_1 - A_0 = \begin{pmatrix}
+                0 & 0 & 0 & 0 \\ 
+                u_{00} & u_{01} & u_{02} & u_{03} \\
+                0 & 0 & 0 & 0 \\ 
+                u_{10} & u_{11} & u_{12} & u_{13} \\
+            \end{pmatrix}
+            = \begin{pmatrix}
+                0 & 0 \\ 1 & 0 \\ 0 & 0 \\ 0 & 1
+            \end{pmatrix}
+            \begin{pmatrix}
+                u_{00} & u_{01} & u_{02} & u_{03} \\
+                u_{10} & u_{11} & u_{12} & u_{13} \\
+            \end{pmatrix}
+            
+        .. code-block:: python
+        
+            u = jnp.array([[u00, u10], [u01, u11], [u02, u12], [u03, u13]])
+            v = jnp.array([1, 3])
+
+        **Simultaneous update of row and column**
+
+        .. math::
+
+            A_1 - A_0 = \begin{pmatrix}
+                0 & 0 & v_0 & 0 \\ 
+                u_0 & u_1 & u_2 + v_1 & u_3 \\
+                0 & 0 & v_2 & 0 \\ 
+                0 & 0 & v_3 & 0 \\ 
+            \end{pmatrix}
+            = \begin{pmatrix}
+                0 & v_0 \\ 1 & v_1 \\ 0 & v_2 \\ 0 & v_3
+            \end{pmatrix}
+            \begin{pmatrix}
+                u_0 & u_1 & u_2 & u_3 \\
+                0 & 0 & 1 & 0 \\
+            \end{pmatrix}
+            
+        .. code-block:: python
+        
+            u = (jnp.array([u0, u1, u2, u3]), 2)
+            v = (jnp.array([v0, v1, v2, v3]), 1)
+    """
+    _check_mat(Ainv)
+    u = _standardize_uv(u, Ainv.shape[0], Ainv.dtype)
+    v = _standardize_uv(v, Ainv.shape[0], Ainv.dtype)
+    _check_uv(u, v)
+
+    R = _get_R(Ainv, u, v)
+    ratio, lufac = _det_and_lufac(R)
+    if return_update:
+        Ainv = _update_Ainv(Ainv, u, v, lufac)
+        return ratio, Ainv
+    else:
+        return ratio
+
+
+class DetCarrier(NamedTuple):
+    Ainv: Array
+    a: Array
+    b: Array
+
+
+def init_det_carrier(A: Array, max_delay: int, max_rank: int = 1) -> DetCarrier:
+    r"""
+    Prepare the data and space for `~lrux.det_lru_delayed`
+
+    :param A:
+        The initial matrix :math:`A_0` with shape (n, n).
+
+    :param max_delay:
+        The maximum iterations T of delayed updates, usually chosen to be ~n/10.
+
+    :param max_rank:
+        The maximum rank K in delayed updates, default to 1.
+
+    :return:
+        A ``NamedTuple`` with the following attributes.
+
+        Ainv:
+            The initial matrix inverse :math:`A_0^{-1}` of shape (n, n).
+        a:
+            The delayed update vectors of shape (T, n, K), initialized to 0
+        b:
+            The delayed update vectors of shape (T, n, K), initialized to 0
+    """
+
+    if max_delay <= 0:
+        raise ValueError(
+            "`max_delay` should be a positive integer. "
+            "Otherwise, please use `det_lru` for non-delayed updates."
+        )
+    _check_mat(A)
+    Ainv = jnp.linalg.inv(A)
+    a = jnp.zeros((max_delay, A.shape[0], max_rank), A.dtype)
+    b = jnp.zeros_like(a)
+    return DetCarrier(Ainv, a, b)
+
+
+def _update_ab(a: Array, new_a: Array, current_delay: int) -> Array:
+    k = new_a.shape[-1]
+    if k > a.shape[-1]:
+        raise ValueError(
+            "The rank of update exceeds max_rank specified in `init_det_carrier`."
+        )
+    return a.at[current_delay, :, :k].set(new_a)
+
+
+def _get_delayed_output(
+    carrier: DetCarrier,
+    u: Tuple[Array, Array],
+    v: Tuple[Array, Array],
+    return_update: bool,
+    current_delay: int,
+) -> Union[Array, Tuple[Array, Array]]:
+    Ainv = carrier.Ainv
+    a = carrier.a[:current_delay]
+    b = carrier.b[:current_delay]
+    R0 = _get_R(Ainv, u, v)
+
+    xuT_a = jnp.einsum("nk,tnl->tkl", u[0], a)
+    euT_a = a[:, u[1], :]
+    uT_a = jnp.concatenate((xuT_a, euT_a), axis=1)
+
+    xvT_b = jnp.einsum("nk,tnl->tkl", v[0], b)
+    evT_b = b[:, v[1], :]
+    vT_b = jnp.concatenate((evT_b, xvT_b), axis=1)
+
+    R = R0 - jnp.einsum("tkl,tml->km", uT_a, vT_b)
+    ratio, lufac = _det_and_lufac(R)
+
+    if return_update:
+        a0 = jnp.concatenate((Ainv[:, v[1]], Ainv @ v[0]), axis=1)
+        new_a = a0 - jnp.einsum("tnk,tlk->nl", a, vT_b)
+        bT0 = jnp.concatenate((u[0].T @ Ainv, Ainv[u[1], :]), axis=0)
+        new_bT = bT0 - jnp.einsum("tkl,tnl->kn", uT_a, b)
+        new_bT = jax.scipy.linalg.lu_solve(lufac, new_bT)
+
+        a = _update_ab(carrier.a, new_a, current_delay)
+        b = _update_ab(carrier.b, new_bT.T, current_delay)
+
+        if current_delay == a.shape[0] - 1:
+            Ainv -= jnp.einsum("tnk,tmk->nm", a, b)
+            carrier = DetCarrier(Ainv, jnp.zeros_like(a), jnp.zeros_like(b))
+        else:
+            carrier = DetCarrier(Ainv, a, b)
+        return ratio, carrier
+    else:
+        return ratio
+
+
+def det_lru_delayed(
+    carrier: DetCarrier,
+    u: Union[ArrayLike, Tuple[Array, ArrayLike]],
+    v: Union[ArrayLike, Tuple[Array, ArrayLike]],
+    return_update: bool = False,
+    current_delay: Optional[int] = None,
+) -> Union[Array, Tuple[Array, DetCarrier]]:
+    r"""
+    Delayed low-rank update of determinant
+
+    :param carrier:
+        The existing delayed update quantities, including :math:`A_0^{-1}`, and
+
+        .. math::
+
+            a_t = A_{t-1}^{-1} v_t
+
+        .. math::
+
+            b_t = (A_{t-1}^{-1})^T u_t
+
+        with :math:`t` from 1 to :math:`\tau-1`. 
+        Initially provided by `~lrux.init_det_carrier`.
+
+    :param u:
+        Low-rank update vector(s) :math:`u_\tau`, the same as :math:`u` in `lrux.det_lru`.
+        The rank of u shouldn't exceed the maximum allowed rank specified 
+        in `~lrux.init_det_carrier`.
+
+    :param v:
+        Low-rank update vector(s) :math:`v_\tau`, the same as :math:`v` in `lrux.det_lru`.
+        The rank of v shouldn't exceed the maximum allowed rank specified 
+        in `~lrux.init_det_carrier`.
+
+    :param return_update:
+        Whether the new carrier with updated quantities should be returned,
+        defaul to False.
+
+    :param current_delay:
+        The current iterations :math:`\tau` of delayed updates. As python starts counting
+        from 0, the actual :math:`\tau` should be ``current_delay + 1``.
+        It must be specified when ``return_update`` is True.
+
+    :return:
+        ratio:
+            The ratio between two determinants
+
+            .. math::
+
+                r_\tau = \frac{\det(A_\tau)}{\det(A_{\tau-1})} = \det(R_\tau)
+
+            where
+
+            .. math::
+
+                R_\tau = I + u_\tau^T A_0^{-1} v_\tau - \sum_{t=1}^{\tau-1} (u_\tau^T a_t) (b_t^T v_\tau)
+
+        new_carrier:
+            Only returned when ``return_update`` is True. The new carrier contains
+            the quantities from the input carrier, and in addition
+
+            .. math::
+
+                a_\tau = A_{\tau-1}^{-1} v_\tau = A_0^{-1} v_\tau - \sum_{t=1}^{\tau-1} a_t (b_t^T v_\tau)
+
+            .. math::
+
+                b_\tau = (A_{\tau-1}^{-1})^T u_\tau = (A_0^{-1})^T u_\tau - \sum_{t=1}^{\tau-1} b_t (a_t^T u_\tau)
+
+            When :math:`\tau` reaches the maximum delayed iterations :math:`T`
+            specified in `~lrux.init_det_carrier`, i.e. ``current_delay == max_delay - 1``, 
+            the current :math:`A_\tau` will be set as the new :math:`A_0`, 
+            whose inverse is given by
+
+            .. math::
+
+                A_\tau^{-1} = A_0^{-1} - \sum_{t=1}^\tau a_t b_t^T
+
+            The ``Ainv`` in ``new_carrier`` will be replaced by :math:`A_\tau^{-1}`, and
+            ``a`` and ``b`` will be set to 0.
+
+    .. warning::
+
+        This function is only recommended for heavy users who understand why and when 
+        to use delayed updates. Otherwise, please choose `~\lrux.det_lru`.
+
+    .. tip::
+
+        Similar to `~lrux.det_lru`, this function is compatible with ``jax.jit`` and 
+        ``jax.vmap``, while ``return_update`` and ``current_delay`` are static arguments
+        which shouldn't be jitted or vmapped.
+
+        We still recommend setting ``donate_argnums=0`` in ``jax.jit`` to reuse 
+        the memory of ``carrier`` if it's no longer needed. For instance,
+
+        .. code-block:: python
+
+            lru_vmap = jax.vmap(det_lru_delayed, in_axes=(0, 0, 0, None, None))
+            lru_jit = jax.jit(lru_vmap, static_argnums=(3, 4), donate_argnums=0)
+
+    Here is a complete example of delayed updates.
+
+    .. code-block:: python
+
+        import os
+        os.environ["JAX_ENABLE_X64"] = "1"
+
+        import random
+        import jax
+        import jax.numpy as jnp
+        import jax.random as jr
+        from lrux import det_lru_delayed, init_det_carrier
+
+        def _get_key():
+            seed = random.randint(0, 2**31 - 1)
+            return jr.key(seed)
+
+        dtype = jnp.float64
+        n = 10
+        max_delay = n // 2
+        max_rank = 2
+        A = jr.normal(_get_key(), (n, n), dtype)
+        carrier = init_det_carrier(A, max_delay, max_rank)
+        detA0 = jnp.linalg.det(A)
+
+        lru_fn = jax.jit(det_lru_delayed, static_argnums=(3, 4), donate_argnums=0)
+
+        for i in range(20):
+            current_delay = i % max_delay
+            k = random.randint(0, max_rank)
+            u = jr.normal(_get_key(), (n, k), dtype)
+            v = jr.normal(_get_key(), (n, k), dtype)
+            ratio, carrier = lru_fn(carrier, u, v, True, current_delay)
+
+            # verify the low-rank update result
+            A += v @ u.T
+            detA1 = jnp.linalg.det(A)
+            assert jnp.allclose(ratio, detA1 / detA0)
+            detA0 = detA1
+    """
+    max_delay = carrier.a.shape[0]
+    if current_delay is None:
+        if return_update:
+            raise ValueError("`current_delay` must be specified to return updates.")
+        current_delay = max_delay - 1
+
+    elif current_delay < 0 or current_delay >= max_delay:
+        raise ValueError(
+            f"`current_delay` should be in range [0, {max_delay}), got {current_delay}."
+        )
+
+    Ainv = carrier.Ainv
+    u = _standardize_uv(u, Ainv.shape[0], Ainv.dtype)
+    v = _standardize_uv(v, Ainv.shape[0], Ainv.dtype)
+    _check_uv(u, v)
+
+    return _get_delayed_output(carrier, u, v, return_update, current_delay)

lrux/pf_lru.py

@@ -0,0 +1,429 @@
+from typing import Optional, Tuple, Union, NamedTuple
+from jax.typing import ArrayLike
+from jax import Array
+import jax
+import jax.numpy as jnp
+from .det_lru import _standardize_uv, _update_ab
+from .pfaffian import skew_eye, pf
+
+
+def _check_mat(mat: Array) -> Array:
+    if mat.ndim != 2 or mat.shape[0] != mat.shape[1] or mat.shape[0] % 2 == 1:
+        raise ValueError(f"Expect input matrix shape (2n, 2n), got {mat.shape}.")
+    return (mat - mat.T) / 2
+
+
+def _get_R(Ainv: Array, u: Tuple[Array, Array]) -> Array:
+    xu_Ainv_xu = jnp.einsum("nk,nm,ml->kl", u[0], Ainv, u[0])
+    xu_Ainv_eu = u[0].T @ Ainv[:, u[1]]
+    eu_Ainv_eu = Ainv[u[1], :][:, u[1]]
+    uT_Ainv_u = jnp.block([[xu_Ainv_xu, xu_Ainv_eu], [-xu_Ainv_eu.T, eu_Ainv_eu]])
+    J = skew_eye(uT_Ainv_u.shape[0] // 2, Ainv.dtype)
+    R = uT_Ainv_u + J
+    return (R - R.T) / 2  # ensure skew-symmetric
+
+
+def _update_Ainv(Ainv: Array, u: Tuple[Array, Array], R: Array) -> Array:
+    Ainv_u = jnp.concatenate((Ainv @ u[0], Ainv[:, u[1]]), axis=1)
+    if R.shape[0] == 2:
+        Ainv_u1, Ainv_u2 = Ainv_u.T
+        outer = jnp.outer(Ainv_u1, Ainv_u2)
+        Ainv -= (outer - outer.T) / R[0, 1]
+    else:
+        Rinv_Ainv_u = jax.scipy.linalg.solve(R, Ainv_u.T)
+        Ainv += Ainv_u @ Rinv_Ainv_u
+    return (Ainv - Ainv.T) / 2
+
+
+def pf_lru(
+    Ainv: Array,
+    u: Union[ArrayLike, Tuple[Array, ArrayLike]],
+    return_update: bool = False,
+) -> Union[Array, Tuple[Array, Array]]:
+    r"""
+    Low-rank update of pfaffian :math:`\mathrm{pf}(A_1) = \mathrm{pf}(A_0 - u J u^T)`.
+    Here :math:`J` is the skew-symmetric identity matrix
+        
+    .. math::
+
+        J = \begin{pmatrix}
+            0 & I \\ -I & 0
+        \end{pmatrix}
+
+    as given in `~lrux.skew_eye`.
+
+    :param Ainv:
+        Inverse of the original skew-symmetric matrix :math:`A_0^{-1}`, shape (n, n)
+
+    :param u:
+        Low-rank update vector(s) :math:`u`, the same as :math:`u` in `lrux.det_lru`.
+
+    :param return_update:
+        Whether the new matrix inverse :math:`A_1^{-1}` should be returned,
+        defaul to False.
+
+    :return:
+        ratio:
+            The ratio between two pfaffians
+
+            .. math::
+
+                r = \frac{\mathrm{pf}(A_1)}{\mathrm{pf}(A_0)} = \frac{\mathrm{pf}(R)}{\mathrm{pf}(J)}
+
+            where
+
+            .. math::
+
+                R = J + u^T A_0^{-1} u
+
+        new_Ainv:
+            The new matrix inverse
+
+            .. math::
+
+                A_1^{-1} = (A_0 + u J u^T)^{-1} = A_0^{-1} + (A_0^{-1} u) R^{-1} (A_0^{-1} u)^T
+
+            Only returned when ``return_update`` is True.
+
+    .. tip::
+
+        This function is compatible with ``jax.jit`` and ``jax.vmap``, while
+        ``return_update`` is a static argument which shouldn't be jitted or vmapped.
+
+        Furthermore, we recommend setting ``donate_argnums=0`` in ``jax.jit`` to reuse 
+        the memory of ``Ainv`` if it's no longer needed. This helps to greatly reduce 
+        the time and memory cost. For instance,
+
+        .. code-block:: python
+
+            lru_vmap = jax.vmap(pf_lru, in_axes=(0, 0, None))
+            lru_jit = jax.jit(lru_vmap, static_argnums=2, donate_argnums=0)
+
+    .. note::
+
+        Here are examples of how to define ``u`` before calling ``pf_lru(Ainv, u)``.
+        Keep in mind that the low-rank update we need should be skew-symmetric and takes
+        the form
+
+        .. math::
+
+            A_1 - A_0 = -u J u^T
+
+        **Update of 1 row and 1 column**
+
+        .. math::
+
+            A_1 - A_0 = \begin{pmatrix}
+                0 & -u_0 & 0 & 0 \\ 
+                u_0 & 0 & u_2 & u_3 \\
+                0 & -u_2 & 0 & 0 \\ 
+                0 & -u_3 & 0 & 0 \\ 
+            \end{pmatrix}
+            = -\begin{pmatrix}
+                u_0 & 0 \\ u_1 & 1 \\ u_2 & 0 \\ u_3 & 0
+            \end{pmatrix}
+            \begin{pmatrix}
+                0 & 1 \\ -1 & 0
+            \end{pmatrix}
+            \begin{pmatrix}
+                u_0 & u_1 & u_2 & u_3 \\
+                0 & 1 & 0 & 0\\
+            \end{pmatrix}
+            
+        .. code-block:: python
+        
+            u = (jnp.array([u0, u1, u2, u3]), 1)
+
+        **Update of 2 rows and 2 columns**
+
+        .. math::
+
+            \begin{split}
+                A_1 - A_0 &= \begin{pmatrix}
+                    0 & -u_{00} & 0 & -u_{10} \\ 
+                    u_{00} & 0 & u_{02} & u_{03} - u_{11} \\
+                    0 & -u_{02} & 0 & -u_{12} \\ 
+                    u_{10} & u_{11} - u_{03} & u_{12} & 0 \\
+                \end{pmatrix} \\
+                &= -\begin{pmatrix}
+                    u_{00} & u_{10} & 0 & 0 \\
+                    u_{01} & u_{11} & 1 & 0 \\
+                    u_{02} & u_{12} & 0 & 0 \\
+                    u_{03} & u_{13} & 0 & 1 \\
+                \end{pmatrix}
+                \begin{pmatrix}
+                    0 & 0 & 1 & 0 \\
+                    0 & 0 & 0 & 1 \\
+                    -1 & 0 & 0 & 0 \\
+                    0 & -1 & 0 & 0 \\
+                \end{pmatrix}
+                \begin{pmatrix}
+                    u_{00} & u_{01} & u_{02} & u_{03} \\
+                    u_{10} & u_{11} & u_{12} & u_{13} \\
+                    0 & 1 & 0 & 0 \\
+                    0 & 0 & 0 & 1 \\
+                \end{pmatrix}
+            \end{split}
+            
+        .. code-block:: python
+        
+            x = jnp.array([[u00, u10], [u01, u11], [u02, u12], [u03, u13]])
+            e = jnp.array([1, 3])
+            u = (x, e)
+    """
+    Ainv = _check_mat(Ainv)
+    u = _standardize_uv(u, Ainv.shape[0], Ainv.dtype)
+    k = u[0].shape[1] + u[1].size
+    if k % 2 == 1:
+        raise ValueError(f"The input u should have even rank, got rank {k}.")
+
+    R = _get_R(Ainv, u)
+    pfR = pf(R)
+    k_half = k // 2
+    ratio = jnp.where((k_half * (k_half - 1) // 2) % 2 == 0, pfR, -pfR)
+    if return_update:
+        Ainv = _update_Ainv(Ainv, u, R)
+        return ratio, Ainv
+    else:
+        return ratio
+
+
+class PfCarrier(NamedTuple):
+    Ainv: Array
+    a: Array
+    Rinv: Array
+
+
+def init_pf_carrier(A: Array, max_delay: int, max_rank: int = 2) -> PfCarrier:
+    r"""
+    Prepare the data and space for `~lrux.pf_lru_delayed`
+
+    :param A:
+        The initial skew-symmetric matrix :math:`A_0` with shape (n, n).
+
+    :param max_delay:
+        The maximum iterations T of delayed updates, usually chosen to be ~n/10.
+
+    :param max_rank:
+        The maximum rank K in delayed updates, default to 2.
+
+    :return:
+        A ``NamedTuple`` with the following attributes.
+
+        Ainv:
+            The initial matrix inverse :math:`A_0^{-1}` of shape (n, n).
+        a:
+            The delayed update vectors of shape (T, n, K), initialized to 0
+        Rinv:
+            The delayed update matrices :math:`R_t^{-1}` of shape (T, K, K), initialized to 0
+    """
+    if max_delay <= 0:
+        raise ValueError(
+            "`max_delay` should be a positive integer. "
+            "Otherwise, please use `pf_lru` for non-delayed updates."
+        )
+    A = _check_mat(A)
+    Ainv = jnp.linalg.inv(A)
+    Ainv = (Ainv - Ainv.T) / 2  # ensure skew-symmetric
+    a = jnp.zeros((max_delay, A.shape[0], max_rank), A.dtype)
+    Rinv = jnp.zeros((max_delay, max_rank, max_rank), A.dtype)
+    return PfCarrier(Ainv, a, Rinv)
+
+
+def _get_delayed_updates(Ainv: Array, a: Array, Rinv: Array) -> Array:
+    if Rinv.shape[-1] == 2:
+        a1 = a[:, :, 0]
+        a2 = a[:, :, 1]
+        outer = jnp.einsum("tn,t,tm->nm", a1, Rinv[:, 0, 1], a2)
+        update = outer - outer.T
+    else:
+        update = jnp.einsum("tnj,tjk,tmk->nm", a, Rinv, a)
+
+    Ainv += update
+    return (Ainv - Ainv.T) / 2  # ensure skew-symmetric
+
+
+def _get_delayed_output(
+    carrier: PfCarrier, u: Tuple[Array, Array], return_update: bool, current_delay: int
+) -> Union[Array, Tuple[Array, Array]]:
+    Ainv = carrier.Ainv
+    a = carrier.a[:current_delay]
+    Rinv = carrier.Rinv[:current_delay]
+    R0 = _get_R(Ainv, u)
+
+    xT_a = jnp.einsum("nk,tnl->tkl", u[0], a)
+    eT_a = a[:, u[1], :]
+    uT_a = jnp.concatenate((xT_a, eT_a), axis=1)
+
+    R = R0 + jnp.einsum("tjk,tkl,tml->jm", uT_a, Rinv, uT_a)
+    pfR = pf(R)
+    k = u[0].shape[1] + u[1].size
+    if k % 2 == 1:
+        raise ValueError(f"The input u should have even rank, got rank {k}.")
+    k_half = k // 2
+    ratio = jnp.where((k_half * (k_half - 1) // 2) % 2 == 0, pfR, -pfR)
+
+    if return_update:
+        a0 = jnp.concatenate((Ainv @ u[0], Ainv[:, u[1]]), axis=1)
+        new_a = a0 + jnp.einsum("tnj,tjk,tlk->nl", a, Rinv, uT_a)
+        a = _update_ab(carrier.a, new_a, current_delay)
+
+        if k == 2:
+            rinv = -1 / ratio
+            new_Rinv = jnp.array([[0, rinv], [-rinv, 0]], dtype=Rinv.dtype)
+        else:
+            new_Rinv = jnp.linalg.inv(R)
+            new_Rinv = (new_Rinv - new_Rinv.T) / 2  # ensure skew-symmetric
+        Rinv = carrier.Rinv.at[current_delay, :k, :k].set(new_Rinv)
+
+        if current_delay == a.shape[0] - 1:
+            Ainv = _get_delayed_updates(Ainv, a, Rinv)
+            carrier = PfCarrier(Ainv, jnp.zeros_like(a), jnp.zeros_like(Rinv))
+        else:
+            carrier = PfCarrier(Ainv, a, Rinv)
+        return ratio, carrier
+    else:
+        return ratio
+
+
+def pf_lru_delayed(
+    carrier: PfCarrier,
+    u: Union[ArrayLike, Tuple[Array, ArrayLike]],
+    return_update: bool = False,
+    current_delay: Optional[int] = None,
+) -> Union[Array, Tuple[Array, PfCarrier]]:
+    r"""
+    Delayed low-rank update of pfaffian
+
+    :param carrier:
+        The existing delayed update quantities, including :math:`A_0^{-1}`, :math:`R_t^{-1}`, and
+
+        .. math::
+
+            a_t = A_{t-1}^{-1} u_t
+
+        with :math:`t` from 1 to :math:`\tau-1`.
+        Initially provided by `~lrux.init_pf_carrier`.
+
+    :param u:
+        Low-rank update vector(s) :math:`u_\tau`, the same as :math:`u` in `lrux.det_lru`.
+        The rank of u shouldn't exceed the maximum allowed rank specified
+        in `~lrux.init_pf_carrier`.
+
+    :param return_update:
+        Whether the new carrier with updated quantities should be returned,
+        defaul to False.
+
+    :param current_delay:
+        The current iterations :math:`\tau` of delayed updates. As python starts counting
+        from 0, the actual :math:`\tau` should be ``current_delay + 1``.
+        It must be specified when ``return_update`` is True.
+
+    :return:
+        ratio:
+            The ratio between two pfaffians
+
+            .. math::
+
+                r_\tau = \frac{\mathrm{pf}(A_\tau)}{\mathrm{pf}(A_{\tau-1})} = \frac{\mathrm{pf}(R_\tau)}{\mathrm{pf}(J)}
+
+            where
+
+            .. math::
+
+                R_\tau = J + u_\tau^T A_0^{-1} u_\tau + \sum_{t=1}^{\tau-1} (u_\tau^T a_t) (a_t^T u_\tau)
+
+        new_carrier:
+            Only returned when ``return_update`` is True. The new carrier contains
+            the quantities from the input carrier, and in addition :math:`R_\tau` and
+
+            .. math::
+
+                a_\tau = A_{\tau-1}^{-1} u_\tau = A_0^{-1} u_\tau + \sum_{t=1}^{\tau-1} a_t R_t^{-1} (a_t^T u_\tau)
+
+            When :math:`\tau` reaches the maximum delayed iterations :math:`T`
+            specified in `~lrux.init_pf_carrier`, i.e. ``current_delay == max_delay - 1``,
+            the current :math:`A_\tau` will be set as the new :math:`A_0`,
+            whose inverse is given by
+
+            .. math::
+
+                A_\tau^{-1} = A_0^{-1} + \sum_{t=1}^\tau a_t R_t^{-1} a_t^T
+
+            The ``Ainv`` in ``new_carrier`` will be replaced by :math:`A_\tau^{-1}`, and
+            ``a`` and ``Rinv`` will be set to 0.
+
+    .. warning::
+
+        This function is only recommended for heavy users who understand why and when 
+        to use delayed updates. Otherwise, please choose `~\lrux.pf_lru`.
+
+    .. tip::
+
+        Similar to `~lrux.det_lru_delayed` and `~lrux.pf_lru`, this function is compatible
+        with ``jax.jit`` and ``jax.vmap``, while ``return_update`` and ``current_delay``
+        are static arguments which shouldn't be jitted or vmapped.
+
+        We still recommend setting ``donate_argnums=0`` in ``jax.jit`` to reuse
+        the memory of ``carrier`` if it's no longer needed. For instance,
+
+        .. code-block:: python
+
+            lru_vmap = jax.vmap(pf_lru_delayed, in_axes=(0, 0, None, None))
+            lru_jit = jax.jit(lru_vmap, static_argnums=(2, 3), donate_argnums=0)
+
+    Here is a complete example of delayed updates.
+
+    .. code-block:: python
+
+        import os
+        os.environ["JAX_ENABLE_X64"] = "1"
+
+        import random
+        import jax
+        import jax.numpy as jnp
+        import jax.random as jr
+        from lrux import skew_eye, pf, init_pf_carrier, pf_lru_delayed
+
+        def _get_key():
+            seed = random.randint(0, 2**31 - 1)
+            return jr.key(seed)
+
+        dtype = jnp.float64
+        n = 10
+        k = 2
+        max_delay = n // 2
+        A = jr.normal(_get_key(), (n, n), dtype)
+        A = (A - A.T) / 2
+        carrier = init_pf_carrier(A, max_delay, k)
+        pfA0 = pf(A)
+
+        lru_fn = jax.jit(pf_lru_delayed, static_argnums=(2, 3), donate_argnums=0)
+
+        for i in range(20):
+            current_delay = i % max_delay
+            ki = random.randint(0, k // 2) * 2  # ensure ki is even
+            u = jr.normal(_get_key(), (n, ki), dtype)
+
+            ratio, carrier = lru_fn(carrier, u, True, current_delay)
+            J = skew_eye(ki // 2, dtype)
+            A -= u @ J @ u.T
+            pfA1 = pf(A)
+            assert jnp.allclose(ratio, pfA1 / pfA0)
+            pfA0 = pfA1
+    """
+    max_delay = carrier.a.shape[0]
+    if current_delay is None:
+        if return_update:
+            raise ValueError("`current_delay` must be specified to return updates.")
+        current_delay = max_delay - 1
+
+    elif current_delay < 0 or current_delay >= max_delay:
+        raise ValueError(
+            f"`current_delay` should be in range [0, {max_delay}), got {current_delay}."
+        )
+
+    Ainv = carrier.Ainv
+    u = _standardize_uv(u, Ainv.shape[0], Ainv.dtype)
+
+    return _get_delayed_output(carrier, u, return_update, current_delay)

lrux/pfaffian.py

@@ -0,0 +1,300 @@
+from typing import Tuple, NamedTuple, Optional, Callable
+from jax import Array
+from functools import partial
+import jax
+import jax.numpy as jnp
+
+
+def skew_eye(n: int, dtype: Optional[jnp.dtype] = None) -> Array:
+    r"""
+    The skew-symmetric identity matrix
+
+    :param n:
+        Number of rows in the output divided by 2.
+
+    :param dtype:
+        Optional dtype, default to floating point.
+
+    :return:
+        The skew-symmetric identity matrix of shape (2n, 2n), defined as
+        
+        .. math::
+
+            J = \begin{pmatrix}
+                0 & I \\ -I & 0
+            \end{pmatrix}
+    """
+    I = jnp.eye(n, dtype=dtype)
+    O = jnp.zeros((n, n), dtype=dtype)
+    return jnp.block([[O, I], [-I, O]])
+
+
+def _pfaffian_direct(A: Array) -> Array:
+    n = A.shape[-1]
+    batch = A.shape[:-2]
+
+    if n % 2 == 1:
+        return jnp.zeros(batch, dtype=A.dtype)
+
+    # By convention, pfaffian of an empty matrix is 1
+    elif n == 0:
+        return jnp.ones(batch, dtype=A.dtype)
+
+    elif n == 2:
+        return A[..., 0, 1]
+
+    elif n == 4:
+        idx = jnp.triu_indices(n, 1)
+        A_upper = A[..., idx[0], idx[1]]
+        a, b, c, d, e, f = jnp.moveaxis(A_upper, -1, 0)
+        return a * f - b * e + d * c
+
+
+def _householder(x: Array, n: Optional[int] = None) -> Tuple[Array, Array, Array]:
+    if n is None:
+        n = 0
+        x0 = x[0]
+        x = x.at[0].set(0)
+    else:
+        x0 = x[n]
+        x = jnp.where(jnp.arange(x.size) <= n, 0, x)
+
+    sigma = jnp.vdot(x, x)
+    norm_x = jnp.sqrt(x0.conj() * x0 + sigma)
+
+    phase = jnp.where(x0 == 0.0, 1.0, jnp.sign(x0))
+    alpha = -phase * norm_x
+
+    v = x.at[n].set(x0 - alpha)
+    v *= jax.lax.rsqrt(jnp.vdot(v, v))
+
+    cond = sigma == 0.0
+    v = jnp.where(cond, 0, v)
+    tau = jnp.where(cond, 0, 2)
+    alpha = jnp.where(cond, x0, alpha)
+
+    return v, tau, alpha
+
+
+@jax.custom_jvp
+def _slogpf_householder(A: Array) -> Array:
+    n = A.shape[0]
+
+    def body_fun(i, val):
+        A, sign, log = val
+        v, tau, alpha = _householder(A[:, i], i + 1)
+        w = tau * A @ v.conj()
+        vw = jnp.outer(v, w)
+        A += vw - vw.T
+
+        new_val = (1 - tau) * jnp.where(i % 2 == 0, -alpha, 1.0)
+        sign *= jnp.sign(new_val)
+        log += jnp.log(jnp.abs(new_val))
+        return A, sign, log
+
+    sign = jnp.array(1, dtype=A.dtype)
+    log = jnp.array(0, dtype=jnp.finfo(A.dtype).dtype)
+    init_val = (A, sign, log)
+    A, sign, log = jax.lax.fori_loop(0, n - 2, body_fun, init_val)
+
+    sign *= jnp.sign(A[n - 2, n - 1])
+    log += jnp.log(jnp.abs(A[n - 2, n - 1]))
+    return sign, log
+
+
+@jax.custom_jvp
+def _slogpf_householder_for(A: Array) -> Tuple[Array, Array]:
+    vals = []
+    for i in range(A.shape[0] - 2):
+        v, tau, alpha = _householder(A[1:, 0])
+        A = A[1:, 1:]
+        w = tau * A @ v.conj()
+        vw = jnp.outer(v, w)
+        A += vw - vw.T
+
+        vals.append(1 - tau)
+        if i % 2 == 0:
+            vals.append(-alpha)
+
+    vals.append(A[-2, -1])
+    vals = jnp.asarray(vals)
+
+    sign = jnp.prod(jnp.sign(vals))
+    log = jnp.sum(jnp.log(jnp.abs(vals)))
+    return sign, log
+
+
+@jax.custom_jvp
+def _slogpf_schur(A: Array) -> Tuple[Array, Array]:
+    T, Z = jax.scipy.linalg.schur(A)
+    vals = jnp.diag(T, k=1)[::2]
+    s, _ = jnp.linalg.slogdet(Z)
+
+    sign = s * jnp.prod(jnp.sign(vals))
+    log = jnp.sum(jnp.log(jnp.abs(vals)))
+    return sign, log
+
+
+def _check_input(
+    A: Array, method: str
+) -> Tuple[Array, Callable[[Array], Tuple[Array, Array]]]:
+    if A.ndim < 2 or A.shape[-2] != A.shape[-1]:
+        raise ValueError(
+            f"The expected input is a square matrix or a batch of them, got input shape {A.shape}."
+        )
+    A = (A - jnp.swapaxes(A, -2, -1)) / 2
+
+    if method == "householder":
+        slogpf_fn = _slogpf_householder
+    elif method == "householder_for":
+        slogpf_fn = _slogpf_householder_for
+    elif method == "schur":
+        if jnp.issubdtype(A, jnp.complexfloating):
+            raise ValueError("The schur method is only available for real dtypes.")
+        slogpf_fn = _slogpf_schur
+    else:
+        raise ValueError(
+            f"Unknown pfaffian method '{method}'. "
+            "Supported methods include 'householder', 'householder_for' and 'schur'."
+        )
+
+    return A, slogpf_fn
+
+
+class SlogpfResult(NamedTuple):
+    sign: Array
+    logabspf: Array
+
+
+@partial(jax.jit, static_argnames=("method",))
+def slogpf(A: Array, *, method: str = "householder") -> SlogpfResult:
+    """
+    Compute the sign and (natural) logarithm of the pfaffian of an array.
+
+    :param A:
+        An array with shape (..., n, n)
+
+    :param method:
+        The method used to compute the pfaffian. Options include
+
+        ``"householder"``:
+            Houserholder transformation, internally using ``jax.lax.fori_loop``
+            to balance between running and compiling time;
+
+        ``"householder_for"``:
+            Houserholder transformation, internally using jitted python for loops
+            to reduce the running time at a cost of much longer compiling time;
+
+        ``"schur""``:
+            Schur decomposition using 
+            `jax.scipy.linalg.schur <https://docs.jax.dev/en/latest/_autosummary/jax.scipy.linalg.schur.html>`_,
+            but only available on CPU and real dtypes.
+
+        The default method is ``"householder"``.
+
+    :return:
+        A ``NamedTuple`` with the following attributes. Both attribtutes have the same
+        batch dimension as the input A.
+
+        ``sign``:
+            sign(A). For a real input, it's 1, 0, or -1. For a complex input, it's a
+            complex number with absolute value 1, or else 0.
+        
+        ``logabspf``:
+            The natural log of the absolute value of the pfaffian.
+    
+    .. tip::
+
+        The input A is always skew-symmetrized.
+
+    .. tip::
+
+        This function has ``jax.custom_jvp`` defined and is backward compatible for both real and complex dtypes.
+    """
+    A, slogpf_fn = _check_input(A, method)
+
+    n = A.shape[-1]
+    if n <= 4 or n % 2 == 1:
+        pfA = _pfaffian_direct(A)
+        return SlogpfResult(jnp.sign(pfA), jnp.log(jnp.abs(pfA)))
+    else:
+        batch = A.shape[:-2]
+        A = A.reshape(-1, n, n)
+
+        outputs = jax.vmap(slogpf_fn)(A)
+        return SlogpfResult(outputs[0].reshape(batch), outputs[1].reshape(batch))
+
+
+@partial(jax.jit, static_argnames=("method",))
+def pf(A: Array, *, method: str = "householder") -> Array:
+    """
+    Compute the pfaffian of an array
+
+    :param A:
+        An array with shape (..., n, n)
+
+    :param method:
+        The method used to compute the pfaffian. Options include
+
+        ``"householder"``:
+            Houserholder transformation, internally using ``jax.lax.fori_loop``
+            to balance between running and compiling time;
+
+        ``"householder_for"``:
+            Houserholder transformation, internally using jitted python for loops
+            to reduce the running time at a cost of much longer compiling time;
+
+        ``"schur""``:
+            Schur decomposition using 
+            `jax.scipy.linalg.schur <https://docs.jax.dev/en/latest/_autosummary/jax.scipy.linalg.schur.html>`_,
+            but only available on CPU and real dtypes.
+
+        The default method is ``"householder"``.
+
+    :return:
+        The pfaffian of A with the same batch dimensions.
+    
+    .. tip::
+
+        The input A is always be skew-symmetrized.
+
+    .. tip::
+
+        This function has ``jax.custom_jvp`` defined and is backward compatible for both real and complex dtypes.
+    """
+    A, slogpf_fn = _check_input(A, method)
+
+    n = A.shape[-1]
+    if n <= 4 or n % 2 == 1:
+        return _pfaffian_direct(A)
+    else:
+        batch = A.shape[:-2]
+        A = A.reshape(-1, n, n)
+
+        sign, log = jax.vmap(slogpf_fn)(A)
+        pfA = sign * jnp.exp(log)
+        return pfA.reshape(batch)
+
+
+def _slogpf_jvp(
+    primals: Tuple[Array], tangents: Tuple[Array], method: str
+) -> Tuple[Tuple[Array, Array], Tuple[Array, Array]]:
+    (A,) = primals
+    (dA,) = tangents
+    A, slogpf_fn = _check_input(A, method)
+
+    sign, ans = slogpf_fn(A)
+    ans_dot = jnp.trace(jnp.linalg.solve(A, dA)) / 2
+
+    if jnp.issubdtype(A.dtype, jnp.complexfloating):
+        sign_dot = 1j * sign * ans_dot.imag
+        ans_dot = ans_dot.real
+    else:
+        sign_dot = jnp.zeros_like(sign)
+
+    return (sign, ans), (sign_dot, ans_dot)
+
+
+_slogpf_householder.defjvp(partial(_slogpf_jvp, method="householder"))
+_slogpf_householder_for.defjvp(partial(_slogpf_jvp, method="householder_for"))
+_slogpf_schur.defjvp(partial(_slogpf_jvp, method="schur"))

lrux-0.1.0.dist-info/METADATA

@@ -0,0 +1,9 @@
+Metadata-Version: 2.4
+Name: lrux
+Version: 0.1.0
+Summary: Fast low-rank updates (LRU) of matrix determinants and pfaffians in JAX
+Author-email: Ao Chen <chenao.phys@gmail.com>, Christopher Roth <christopher_roth@utexas.edu>
+Requires-Python: >=3.8
+License-File: LICENSE
+Requires-Dist: jax>=0.4.4
+Dynamic: license-file

lrux-0.1.0.dist-info/RECORD

@@ -0,0 +1,9 @@
+lrux/__init__.py,sha256=nD_NrlBN23t7aiwiSdsuHEK8f4FJwt9t5Bmmqmrr1hA,190
+lrux/det_lru.py,sha256=6TKJkZnErS8lLlYOc1Tnr-TayCJ7p0lzujGSssbO9fs,17630
+lrux/pf_lru.py,sha256=K53igWzN8Q4Kq93hsS_TKCI7o8xOsXleRkRzegkOEw4,14036
+lrux/pfaffian.py,sha256=_QaABw8kBB9Yy1YugdQJzdtppreP0wbSEqwvCS57otU,8590
+lrux-0.1.0.dist-info/licenses/LICENSE,sha256=XK1r3QjKtOz1L8u1S5X9RLysQI1z7JWAvwf6mEgjvco,1064
+lrux-0.1.0.dist-info/METADATA,sha256=HfVS0qJiXgjI9wfUElCwXCdu_rNJrw1QRDwOLG6dVyg,316
+lrux-0.1.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+lrux-0.1.0.dist-info/top_level.txt,sha256=O8y6IfUNFFU5hU-DKZ6WGcK-AIh--pAvaoE6gxzkDyc,5
+lrux-0.1.0.dist-info/RECORD,,

lrux-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (80.9.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

lrux-0.1.0.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Chen Ao
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

lrux-0.1.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+lrux