folx 0.2.3__tar.gz → 0.2.4__tar.gz

{folx-0.2.3 → folx-0.2.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: folx
-Version: 0.2.3
+Version: 0.2.4
 Summary: Forward Laplacian for JAX
 Home-page: https://github.com/microsoft/folx
 License: MIT
@@ -28,6 +28,7 @@ Requires-Dist: jax
 Requires-Dist: jaxlib
 Requires-Dist: jaxtyping
 Requires-Dist: numpy
+Requires-Dist: parameterized
 Requires-Dist: pytest
 Project-URL: Repository, https://github.com/microsoft/folx
 Description-Content-Type: text/markdown

folx-0.2.4/folx/ad.py

@@ -0,0 +1,104 @@
+import jax
+import jax.numpy as jnp
+import jax.flatten_util as jfu
+import jax.tree_util as jtu
+
+
+def is_tree_complex(tree):
+    leaves = jtu.tree_leaves(tree)
+    return any(jnp.iscomplexobj(leaf) for leaf in leaves)
+
+
+def vjp_rc(fun, *primals: jax.Array):
+    def real_fun(*primals):
+        return jnp.real(fun(*primals))
+
+    def imag_fun(*primals):
+        return jnp.imag(fun(*primals))
+
+    _, vjp_r = jax.vjp(real_fun, *primals)
+    _, vjp_i = jax.vjp(imag_fun, *primals)
+
+    def vjp(*tangents: jax.Array):
+        real_tangents = jtu.tree_map(jnp.real, tangents)
+        imag_tangents = jtu.tree_map(jnp.imag, tangents)
+
+        # letters: v=vector, j=jacobian, r=real, i=imag
+        vr_jr = vjp_r(*real_tangents)
+        vi_jr = vjp_r(*imag_tangents)
+        vr_ji = vjp_i(*real_tangents)
+        vi_ji = vjp_i(*imag_tangents)
+
+        result = jtu.tree_map(
+            lambda vr_jr, vi_jr, vr_ji, vi_ji: vr_jr - vi_ji + 1j * (vr_ji + vi_jr),
+            vr_jr,
+            vi_jr,
+            vr_ji,
+            vi_ji,
+        )
+        return result
+
+    return vjp
+
+
+def vjp(fun, *primals: jax.Array):
+    out, vjp = jax.vjp(fun, *primals)
+    if is_tree_complex(primals):
+        if is_tree_complex(out):
+            # C -> C
+            return vjp
+        else:
+            # C -> R
+            return vjp
+    else:
+        if is_tree_complex(out):
+            # R -> C
+            return vjp_rc(fun, *primals)
+        else:
+            # R -> R
+            return vjp
+
+
+def jacrev(f):
+    # Similar to jax.jacrev but works with complex inputs and outputs.
+    # A crucial difference is that we do not preserve the structure of the output but
+    # always flatten it to a 1D array.
+    def jacfun(*primals):
+        flat_primals, unravel = jfu.ravel_pytree(primals)
+
+        def flat_f(x):
+            return jfu.ravel_pytree(f(*unravel(x)))[0]
+
+        out = flat_f(flat_primals)
+
+        result = jax.vmap(vjp(flat_f, flat_primals))(
+            jnp.eye(out.size, dtype=out.dtype)
+        )[0]
+        result = jax.vmap(unravel, out_axes=0)(result)
+        if len(primals) == 1:
+            return result[0]
+        return result
+
+    return jacfun
+
+
+def jacfwd(f):
+    # Similar to jax.jacfwd but works with complex inputs and outputs.
+    # A crucial difference is that we do not preserve the structure of the input but
+    # always flatten it to a 1D array.
+    def jacfun(*primals):
+        flat_primals, unravel = jfu.ravel_pytree(primals)
+
+        def jvp_fun(s):
+            return jax.jvp(f, primals, unravel(s))[1]
+
+        eye = jnp.eye(flat_primals.size, dtype=flat_primals.dtype)
+        J = jax.vmap(jvp_fun, out_axes=-1)(eye)
+        return J
+
+    return jacfun
+
+
+def hessian(f):
+    # Similar to jax.hessian but works with complex inputs and outputs.
+    return jacfwd(jacrev(f))

{folx-0.2.3 → folx-0.2.4}/folx/api.py

@@ -151,7 +151,7 @@ class FwdJacobian(NamedTuple):
         if idx is None:
             return self.dense_array
         if len(idx) == 0:
-            return jnp.zeros((*self.data_shape, 0))
+            return jnp.zeros((*self.data_shape, 0), dtype=self.data.dtype)
         if self.x0_idx is None:
             return self.data[idx]
         return self.materialize_for_idx(self.get_index_mask(idx), len(idx))
@@ -225,6 +225,17 @@ class FwdJacobian(NamedTuple):
     def astype(self, dtype):
         return FwdJacobian(self.data.astype(dtype), self.x0_idx)
 
+    @property
+    def real(self):
+        return FwdJacobian(self.data.real, self.x0_idx)
+
+    @property
+    def imag(self):
+        return FwdJacobian(self.data.imag, self.x0_idx)
+
+    def conj(self):
+        return FwdJacobian(self.data.conj(), self.x0_idx)
+
 
 class FwdLaplArray(NamedTuple):
     """
@@ -265,12 +276,8 @@ class FwdLaplArray(NamedTuple):
         return FwdLaplArray(self.x, self.jacobian.as_dense, self.laplacian)
 
     def astype(self, dtype):
-        if dtype in (
-            jnp.float16,
-            jnp.float32,
-            jnp.float64,
-            jnp.complex64,
-            jnp.complex128,
+        if jax.dtypes.issubdtype(dtype, jnp.floating) or jax.dtypes.issubdtype(
+            dtype, jnp.complexfloating
         ):
             return FwdLaplArray(
                 self.x.astype(dtype),
@@ -280,6 +287,21 @@ class FwdLaplArray(NamedTuple):
         # If we convert to integer or boolean we drop the derivatives
         return self.x.astype(dtype)
 
+    @property
+    def dtype(self):
+        return self.x.dtype
+
+    @property
+    def real(self):
+        return FwdLaplArray(self.x.real, self.jacobian.real, self.laplacian.real)
+
+    @property
+    def imag(self):
+        return FwdLaplArray(self.x.imag, self.jacobian.imag, self.laplacian.imag)
+
+    def conj(self):
+        return FwdLaplArray(self.x.conj(), self.jacobian.conj(), self.laplacian.conj())
+
 
 def IS_LPL_ARR(x):
     return isinstance(x, FwdLaplArray)

{folx-0.2.3 → folx-0.2.4}/folx/custom_hessian.py

@@ -1,6 +1,8 @@
 import jax
 import jax.numpy as jnp
 
+from folx.ad import is_tree_complex
+
 from .api import Array, ExtraArgs, FwdLaplArgs, MergeFn, JAC_DIM
 from .utils import trace_of_product
 
@@ -56,4 +58,9 @@ def slogdet_jac_hessian_jac(
         return None, elementwise(*x)
 
     signs, flat_out = jax.lax.scan(scan_wrapper, None, (A_inv, J))[1]
-    return signs.reshape(leading_dims), flat_out.reshape(leading_dims)
+    sign_out, log_abs_out = signs.reshape(leading_dims), flat_out.reshape(leading_dims)
+
+    if is_tree_complex(A):
+        # this is not the real Tr(JHJ^T) but a cached value we use later to compute the Tr(JHJ^T)
+        return log_abs_out, log_abs_out.real
+    return sign_out, log_abs_out.real

{folx-0.2.3 → folx-0.2.4}/folx/hessian.py

@@ -34,6 +34,7 @@ from .utils import (
     trace_of_product,
     vmap_sequences_and_squeeze,
 )
+from .ad import hessian, jacrev
 
 
 def general_jac_hessian_jac(
@@ -56,8 +57,8 @@ def general_jac_hessian_jac(
     grad_2d = jnp.concatenate([x.T for x in grads_2d], axis=0)
     D, K = grad_2d.shape
     if K > D:
-        # jax.hessian uses Fwd on Reverse AD
-        flat_hessian = jax.hessian(flat_fn)(flat_x)
+        # Fwd on Reverse AD
+        flat_hessian = hessian(flat_fn)(flat_x)
         flat_out = trace_of_product(flat_hessian, grad_2d @ grad_2d.T)
     elif D > K:
         # Directly copmute the trace of tr(HJJ^T)=tr(J^THJ)
@@ -79,7 +80,7 @@ def general_jac_hessian_jac(
         @functools.partial(jax.vmap, in_axes=-1, out_axes=-1)
         def hvp(tangent):
             def jacobian(x):
-                return jax.jacrev(flat_fn)(x)
+                return jacrev(flat_fn)(x)
 
             return jax.jvp(jacobian, (flat_x,), (tangent,))[1]
 
@@ -115,7 +116,9 @@ def off_diag_jac_hessian_jac(
     return unravel(flat_out)
 
 
-def mul_jac_hessian_jac(fn: ForwardFn, args: FwdLaplArgs, shared_idx: Array | None):
+def dot_product_jac_hessian_jac(
+    fn: ForwardFn, args: FwdLaplArgs, shared_idx: Array | None
+):
     # For a dot product we know that the hessian looks like this:
     # [0, I]
     # [I, 0]
@@ -321,7 +324,7 @@ def vmapped_jac_hessian_jac(
         if custom_jac_hessian_jac is not None:
             result = custom_jac_hessian_jac(args, extra_args, merge, materialize_idx)
         elif FunctionFlags.MULTIPLICATION in flags:
-            result = mul_jac_hessian_jac(merged_fn, args, materialize_idx)
+            result = dot_product_jac_hessian_jac(merged_fn, args, materialize_idx)
         elif FunctionFlags.LINEAR_IN_ONE in flags:
             result = off_diag_jac_hessian_jac(merged_fn, args, materialize_idx)
         else:

{folx-0.2.3 → folx-0.2.4}/folx/interpreter.py

@@ -161,8 +161,8 @@ def eval_jaxpr_with_forward_laplacian(
 
     def eval_laplacian(eqn: core.JaxprEqn, invals):
         subfuns, params = eqn.primitive.get_bind_params(eqn.params)
-        fn = get_laplacian(eqn.primitive, True)
         with LoggingPrefix(f'({summarize(eqn.source_info)})'):
+            fn = get_laplacian(eqn.primitive, True)
             return fn(
                 (*subfuns, *invals), params, sparsity_threshold=sparsity_threshold
             )
@@ -219,12 +219,18 @@ def init_forward_laplacian_state(
     x_flat, unravel = ravel(x)
     jac = jtu.tree_map(jnp.ones_like, x)
     jac_idx = unravel(np.arange(x_flat.shape[0]))
+    if jnp.iscomplexobj(x_flat):
+        logging.info(
+            '[folx] Found complex input. This is not well supported, results might be wrong.'
+        )
     if sparsity:
         jac = jtu.tree_map(
-            lambda j, i: FwdJacobian(j[None], np.array(i)[None]), jac, jac_idx
+            lambda j, i: FwdJacobian(j.astype(x_flat.dtype)[None], np.array(i)[None]),
+            jac,
+            jac_idx,
         )
     else:
-        jac = jax.vmap(unravel)(jnp.eye(len(x_flat)))
+        jac = jax.vmap(unravel)(jnp.eye(len(x_flat), dtype=x_flat.dtype))
         jac = jtu.tree_map(FwdJacobian.from_dense, jac)
     lapl_x = jtu.tree_map(jnp.zeros_like, x)
     return jtu.tree_map(FwdLaplArray, x, jac, lapl_x)

{folx-0.2.3 → folx-0.2.4}/folx/jvp.py

@@ -29,6 +29,7 @@ from .utils import (
     get_jacobian_for_reduction,
     np_concatenate_brdcast,
 )
+from .ad import vjp
 
 R = TypeVar('R', bound=PyTree[Array])
 
@@ -115,7 +116,7 @@ def sparse_diag_jvp(
     ):
         # If we have elementwise functions, we can just compute the full jacobian and
         # do the operations a bit faster.
-        jac = jax.grad(lambda x: jnp.sum(fwd(x)))(laplace_args.x[0])  # type: ignore
+        jac = vjp(fwd, laplace_args.x[0])(jnp.ones_like(y))[0]
         grad_y = jac * laplace_args.jacobian[0].data
         lapl_y = jac * laplace_args.laplacian[0]
     else:
@@ -373,7 +374,7 @@ def dense_elementwise_jvp(
     if y.shape != laplace_args.x[0].shape:
         return dense_split_jvp(fwd, laplace_args)
 
-    jac = jax.grad(lambda x: jnp.sum(fwd(x)))(laplace_args.x[0])  # type: ignore
+    jac = vjp(fwd, laplace_args.x[0])(jnp.ones_like(y))[0]
     grad_y = jac * laplace_args.dense_jacobian[0]
     lapl_y = jac * laplace_args.laplacian[0]
     return y, grad_y, lapl_y

{folx-0.2.3 → folx-0.2.4}/folx/wrapped_functions.py

@@ -7,16 +7,23 @@ import jax.numpy as jnp
 import jax.tree_util as jtu
 from jax.core import Primitive
 
+from folx.ad import is_tree_complex
+
 from .api import (
     Array,
     ArrayOrFwdLaplArray,
     ForwardLaplacian,
     FunctionFlags,
+    FwdJacobian,
     FwdLaplArray,
     PyTree,
 )
 from .custom_hessian import slogdet_jac_hessian_jac
-from .wrapper import wrap_forward_laplacian, warp_without_fwd_laplacian
+from .wrapper import (
+    wrap_elementwise,
+    wrap_forward_laplacian,
+    warp_without_fwd_laplacian,
+)
 
 R = TypeVar('R', bound=PyTree[Array])
 P = ParamSpec('P')
@@ -119,7 +126,10 @@ def dtype_conversion(
 def slogdet(x):
     # We only need this custom slog det to avoid a jax bug
     # https://github.com/google/jax/issues/17379
-    return jnp.linalg.slogdet(x)
+    # We explictily decompose this here as newer version will return
+    # SlogDetResult which is a NamedTuple and does not combine nicely with regular tuples.
+    sign, logdet = jnp.linalg.slogdet(x)
+    return sign, logdet
 
 
 def slogdet_jvp(primals, tangents):
@@ -136,10 +146,18 @@ def slogdet_jvp(primals, tangents):
 
     jacobians = jnp.linalg.inv(primals)
 
-    def custom_jvp(jacobian, tangent):
-        return (jnp.zeros(()), jnp.vdot(jacobian.T, tangent))
+    def custom_jvp(jacobian, tangent, sign):
+        jac_dot_tangent = jnp.vdot(jacobian.T.conj(), tangent)
+        if jac_dot_tangent.dtype in (jnp.complex64, jnp.complex128):
+            # this is not the real jvp but a cached value to ease the Tr(JHJ^T) computation
+            sign_jvp = jac_dot_tangent
+            log_det_jvp = jac_dot_tangent.real
+        else:
+            sign_jvp = jnp.zeros(())
+            log_det_jvp = jac_dot_tangent
+        return (sign_jvp, log_det_jvp)
 
-    y_tangent = jax.vmap(custom_jvp)(jacobians, tangents)
+    y_tangent = jax.vmap(custom_jvp)(jacobians, tangents, sign)
 
     y, y_tangent = jtu.tree_map(lambda x: x.reshape(*batch_shape), (y, y_tangent))
     return y, y_tangent
@@ -158,15 +176,44 @@ def slogdet_wrapper(
     )
     sign, logdet = fwd_lapl_fn(x, {}, sparsity_threshold=0)
     # Remove the jacobian of the sign
-    sign = warp_without_fwd_laplacian(lambda x: x)((sign,), {}, sparsity_threshold=0)
+    if jax.dtypes.issubdtype(sign.dtype, jnp.complexfloating):
+        sign_jac = sign.jacobian.data
+        sign_jac_flat = sign_jac.reshape(-1, *sign.shape).imag
+        sign_jac_dot = jnp.einsum('i...,i...->...', sign_jac_flat, sign_jac_flat)
+        sign = FwdLaplArray(
+            sign.x,
+            FwdJacobian(1.0j * sign.x * sign_jac.imag, x0_idx=sign.jacobian.x0_idx),
+            sign.x * (1.0j * sign.laplacian.imag - sign_jac_dot),
+        )
+    if jax.dtypes.issubdtype(sign.dtype, jnp.floating):
+        sign = sign.x
     return sign, logdet
 
 
+def abs_wrapper(
+    x: tuple[ArrayOrFwdLaplArray],
+    kwargs: dict[str, Any],
+    sparsity_threshold: int,
+):
+    if not is_tree_complex(x):
+        return wrap_forward_laplacian(
+            jax.lax.abs, flags=FunctionFlags.LINEAR, in_axes=()
+        )(x, kwargs, sparsity_threshold)
+    import folx  # we must import folx here to avoid circular imports
+
+    return folx.forward_laplacian(
+        lambda x: jnp.sqrt((x * x.conj()).real),
+        sparsity_threshold=sparsity_threshold,
+        disable_jit=True,
+    )(x[0])
+
+
 _LAPLACE_FN_REGISTRY: dict[Primitive | str, ForwardLaplacian] = {
+    jax.lax.conj_p: wrap_elementwise(jnp.conj),
+    jax.lax.imag_p: wrap_elementwise(jnp.imag),
+    jax.lax.real_p: wrap_elementwise(jnp.real),
     jax.lax.dot_general_p: dot_general,
-    jax.lax.abs_p: wrap_forward_laplacian(
-        jax.lax.abs, flags=FunctionFlags.LINEAR, in_axes=()
-    ),
+    jax.lax.abs_p: abs_wrapper,
     jax.lax.neg_p: wrap_forward_laplacian(
         jax.lax.neg, flags=FunctionFlags.LINEAR, in_axes=()
     ),
@@ -266,11 +313,11 @@ _LAPLACE_FN_REGISTRY: dict[Primitive | str, ForwardLaplacian] = {
         flags=FunctionFlags.INDEXING | FunctionFlags.SCATTER,
         name='scatter',
     ),
-    jax.lax.scatter_add_p: wrap_forward_laplacian(
-        jax.lax.scatter_add_p.bind,
-        flags=FunctionFlags.LINEAR | FunctionFlags.SCATTER,
-        name='scatter_add',
-    ),
+    # jax.lax.scatter_add_p: wrap_forward_laplacian(
+    #     jax.lax.scatter_add_p.bind,
+    #     flags=FunctionFlags.LINEAR | FunctionFlags.SCATTER,
+    #     name='scatter_add',
+    # ),
     jax.lax.stop_gradient_p: warp_without_fwd_laplacian(jax.lax.stop_gradient),
     jax.lax.eq_p: warp_without_fwd_laplacian(jax.lax.eq),
     jax.lax.lt_p: warp_without_fwd_laplacian(jax.lax.lt),

{folx-0.2.3 → folx-0.2.4}/folx/wrapper.py

@@ -14,6 +14,7 @@ from .api import (
     ForwardLaplacian,
     ForwardLaplacianFns,
     FunctionFlags,
+    FwdJacobian,
     FwdLaplArgs,
     FwdLaplArray,
     MergeFn,
@@ -120,7 +121,12 @@ def wrap_forward_laplacian(
         lapl_y = tree_add(
             lapl_y, lapl_fns.jac_hessian_jac_trace(laplace_args, sparsity_threshold)
         )
-        return jax.tree_util.tree_map(FwdLaplArray, y, grad_y, lapl_y)
+        return jax.tree_util.tree_map(
+            lambda x, jac, lapl: FwdLaplArray(x, jac, lapl).astype(x.dtype),
+            y,
+            grad_y,
+            lapl_y,
+        )
 
     return new_fn
 
@@ -139,3 +145,25 @@ def warp_without_fwd_laplacian(fn) -> ForwardLaplacian:
         return fn(*args, **kwargs)
 
     return wrapped
+
+
+def wrap_elementwise(fn) -> ForwardLaplacian:
+    """
+    Decorator that applies the function directly to the data, jacobian and laplacian.
+    """
+
+    def wrapped(args, kwargs, sparsity_threshold: int):
+        assert len(args) == 1
+
+        def inner(x: ArrayOrFwdLaplArray):
+            if isinstance(x, FwdLaplArray):
+                return FwdLaplArray(
+                    fn(x.x),
+                    FwdJacobian(fn(x.jacobian.data), x.jacobian.x0_idx),
+                    fn(x.laplacian),
+                )
+            return fn(x)
+
+        return inner(args[0])
+
+    return wrapped

{folx-0.2.3 → folx-0.2.4}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = 'folx'
-version = '0.2.3'
+version = '0.2.4'
 description = 'Forward Laplacian for JAX'
 authors = [
     "Nicholas Gao <n.gao@tum.de>",
@@ -37,6 +37,7 @@ jaxlib = "*"
 jaxtyping = "*"
 numpy = "*"
 pytest = "*"
+parameterized = "*"
 
 [build-system]
 requires = ["poetry-core>=1.0.0"]