Add JaxLinearOperator

src/operators/jax_operator.py

@@ -23,7 +23,7 @@ from .operator import Operator
 try:
     import jax
     jax.config.update("jax_enable_x64", True)
-    __all__ = ["JaxOperator", "JaxLikelihoodEnergyOperator"]
+    __all__ = ["JaxOperator", "JaxLikelihoodEnergyOperator", "JaxLinearOperator"]
 except ImportError:
     __all__ = []
 
@@ -65,7 +65,7 @@ class JaxOperator(Operator):
         if is_linearization(x):
             res, bwd = self._vjp(x.val.val)
             fwd = lambda y: self._fwd(x.val.val, y)
-            jac = _JaxJacobian(self._domain, self._target, fwd, bwd)
+            jac = _JaxLinearOperator(self._domain, self._target, fwd, bwd)
             return x.new(makeField(self._target, _jax2np(res)), jac)
         res = _jax2np(self._func(x.val))
         if isinstance(res, dict):
@@ -100,6 +100,43 @@ class JaxOperator(Operator):
         return None, JaxOperator(dom, self._target, func2)
 
 
+def JaxLinearOperator(domain, target, func, domain_dtype):
+    """Wrap a jax function as nifty linear operator.
+
+    Parameters
+    ----------
+    domain : DomainTuple or MultiDomain
+        Domain of the operator.
+
+    target : DomainTuple or MultiDomain
+        Target of the operator.
+
+    func : callable
+        The jax function that is evaluated by the operator. It has to be
+        implemented in terms of `jax.numpy` calls. If `domain` is a
+        `DomainTuple`, `func` takes a `dict` as argument and like-wise for the
+        target.
+
+    domain_dtype:
+        Dtype of the domain. If `domain` is a `MultiDomain`, `domain_dtype` is
+        supposed to be a dictionary.
+
+    Note
+    ----
+    It is the user's responsibility that func is actually a linear function. The
+    user can double check this with the help of `nifty8.extra.check_linear_operator`.
+    """
+    from ..domain_tuple import DomainTuple
+    from ..sugar import makeDomain
+    domain = makeDomain(domain)
+    if isinstance(domain, DomainTuple):
+        inp = np.ones(domain.shape, domain_dtype)
+    else:
+        inp = {kk: np.ones(domain[kk].shape, domain_dtype[kk]) for kk in domain.keys()}
+    func_transposed = jax.jit(jax.vjp(func, inp)[1])
+    return _JaxLinearOperator(domain, target, func, func_transposed)
+
+
 class JaxLikelihoodEnergyOperator(LikelihoodEnergyOperator):
     """Wrap a jax function as nifty likelihood energy operator.
 
@@ -163,7 +200,7 @@ class JaxLikelihoodEnergyOperator(LikelihoodEnergyOperator):
         return None, JaxLikelihoodEnergyOperator(dom, func2, trafo, dt)
 
 
-class _JaxJacobian(LinearOperator):
+class _JaxLinearOperator(LinearOperator):
     def __init__(self, domain, target, func, func_transposed):
         from ..sugar import makeDomain
         self._domain = makeDomain(domain)

test/test_operators/test_jax.py

@@ -29,16 +29,26 @@ pmp = pytest.mark.parametrize
 
 
 @pmp("dom", [ift.RGSpace((10, 8)), (ift.RGSpace(10), ift.RGSpace(8))])
-@pmp("func", [lambda x: x, lambda x: x**2, lambda x: x*x, lambda x: x*x[0, 0],
-              lambda x: jnp.sin(x), lambda x: x*x.sum()])
+@pmp("func", [(lambda x: x, True), (lambda x: x**2, False), (lambda x: x*x, False),
+              (lambda x: x*x[0, 0], False), (lambda x: x+x[0, 0], True),
+              (lambda x: jnp.sin(x), False), (lambda x: x*x.sum(), False),
+              (lambda x: x+x.sum(), True)])
 def test_jax(dom, func):
     pytest.importorskip("jax")
     loc = ift.from_random(dom)
-    res0 = np.array(func(loc.val))
-    op = ift.JaxOperator(dom, dom, func)
+    f, linear = func
+    res0 = np.array(f(loc.val))
+    op = ift.JaxOperator(dom, dom, f)
     np.testing.assert_allclose(res0, op(loc).val)
     ift.extra.check_operator(op, ift.from_random(op.domain))
 
+    op = ift.JaxLinearOperator(dom, dom, f, np.float64)
+    if linear:
+        ift.extra.check_linear_operator(op)
+    else:
+        with pytest.raises(AssertionError):
+            ift.extra.check_linear_operator(op)
+
 
 def test_mf_jax():
     pytest.importorskip("jax")