Removed AuxLinearOperator.

This operator was originally introduced as a hack for Optimistix. It's now been unused for a while.

Author: patrick-kidger

lineax/__init__.py

@@ -18,7 +18,6 @@
 from ._operator import (
     AbstractLinearOperator as AbstractLinearOperator,
     AddLinearOperator as AddLinearOperator,
-    AuxLinearOperator as AuxLinearOperator,
     ComposedLinearOperator as ComposedLinearOperator,
     conj as conj,
     diagonal as diagonal,

lineax/_misc.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from collections.abc import Callable
 
 import equinox as eqx
 import jax
@@ -39,13 +38,6 @@ def resolve_rcond(rcond, n, m, dtype):
         return jnp.where(rcond < 0, jnp.finfo(dtype).eps, rcond)
 
 
-class NoneAux(eqx.Module):
-    fn: Callable
-
-    def __call__(self, *args, **kwargs):
-        return self.fn(*args, **kwargs), None
-
-
 def jacobian(fn, in_size, out_size, holomorphic=False, has_aux=False, jac=None):
     if jac is None:
         # Heuristic for which is better in each case

lineax/_operator.py

@@ -43,7 +43,6 @@
     default_floating_dtype,
     inexact_asarray,
     jacobian,
-    NoneAux,
     strip_weak_dtype,
 )
 from ._tags import (
@@ -524,14 +523,6 @@ def __call__(self, x):
         return self.fn(x, self.args)
 
 
-class _NoAuxOut(eqx.Module):
-    fn: Callable
-
-    def __call__(self, x):
-        f, _ = self.fn(x)
-        return f
-
-
 class _Unwrap(eqx.Module):
     fn: Callable
 
@@ -573,16 +564,15 @@ class JacobianLinearOperator(AbstractLinearOperator):
     tags: frozenset[object] = eqx.field(static=True)
     jac: Literal["fwd", "bwd"] | None
 
-    @eqxi.doc_remove_args("closure_convert", "_has_aux")
+    @eqxi.doc_remove_args("closure_convert")
     def __init__(
         self,
         fn: Callable,
         x: PyTree[ArrayLike],
         args: PyTree[Any] = None,
         tags: object | Iterable[object] = (),
-        closure_convert: bool = True,
-        _has_aux: bool = False,  # TODO(kidger): remove, no longer used
         jac: Literal["fwd", "bwd"] | None = None,
+        closure_convert: bool = True,
     ):
         """**Arguments:**
 
@@ -606,8 +596,6 @@ def __init__(
                 "`jac` argument of `JacobianLinearOperator` should be either "
                 "`'fwd'`, `'bwd'`, or `None`."
             )
-        if not _has_aux:
-            fn = NoneAux(fn)
         # Flush out any closed-over values, so that we can safely pass `self`
         # across API boundaries. (In particular, across `linear_solve_p`.)
         # We don't use `jax.closure_convert` as that only flushes autodiffable
@@ -625,7 +613,7 @@ def __init__(
         self.jac = jac
 
     def mv(self, vector):
-        fn = _NoAuxOut(_NoAuxIn(self.fn, self.args))
+        fn = _NoAuxIn(self.fn, self.args)
         if self.jac == "fwd" or self.jac is None:
             _, out = jax.jvp(fn, (self.x,), (vector,))
         elif self.jac == "bwd":
@@ -651,7 +639,7 @@ def as_matrix(self):
     def transpose(self):
         if is_symmetric(self):
             return self
-        fn = _NoAuxOut(_NoAuxIn(self.fn, self.args))
+        fn = _NoAuxIn(self.fn, self.args)
         # Works because vjpfn is a PyTree
         _, vjpfn = jax.vjp(fn, self.x)
         vjpfn = _Unwrap(vjpfn)
@@ -663,7 +651,7 @@ def in_structure(self):
         return strip_weak_dtype(jax.eval_shape(lambda: self.x))
 
     def out_structure(self):
-        fn = _NoAuxOut(_NoAuxIn(self.fn, self.args))
+        fn = _NoAuxIn(self.fn, self.args)
         return strip_weak_dtype(eqxi.cached_filter_eval_shape(fn, self.x))
 
 
@@ -1168,30 +1156,6 @@ def out_structure(self):
         return self.operator1.out_structure()
 
 
-class AuxLinearOperator(AbstractLinearOperator):
-    """Internal to lineax. Used to represent a linear operator with additional
-    metadata attached.
-    """
-
-    operator: AbstractLinearOperator
-    aux: PyTree[Array]
-
-    def mv(self, vector):
-        return self.operator.mv(vector)
-
-    def as_matrix(self):
-        return self.operator.as_matrix()
-
-    def transpose(self):
-        return self.operator.transpose()
-
-    def in_structure(self):
-        return self.operator.in_structure()
-
-    def out_structure(self):
-        return self.operator.out_structure()
-
-
 #
 # Operations on `AbstractLinearOperator`s.
 # These are done through `singledispatch` rather than as methods.
@@ -1260,7 +1224,7 @@ def _(operator):
     if operator.jac == "bwd":
         # For backward mode, use VJP + linear_transpose.
         # This works even with custom_vjp functions that don't support forward-mode AD.
-        _, vjp_fn, aux = jax.vjp(fn, operator.x, has_aux=True)
+        _, vjp_fn = jax.vjp(fn, operator.x)
         if is_symmetric(operator):
             # For symmetric: J = J.T, so vjp directly gives J @ v
             lin = _Unwrap(vjp_fn)
@@ -1270,10 +1234,8 @@ def _(operator):
                 jax.linear_transpose(lambda g: vjp_fn(g)[0], operator.out_structure())
             )
     else:  # "fwd" or None
-        (_, aux), lin = jax.linearize(fn, operator.x)
-        lin = _NoAuxOut(lin)
-    out = FunctionLinearOperator(lin, operator.in_structure(), operator.tags)
-    return AuxLinearOperator(out, aux)
+        _, lin = jax.linearize(fn, operator.x)
+    return FunctionLinearOperator(lin, operator.in_structure(), operator.tags)
 
 
 # materialise
@@ -1346,16 +1308,14 @@ def _(operator):
 @materialise.register(JacobianLinearOperator)
 def _(operator):
     fn = _NoAuxIn(operator.fn, operator.args)
-    jac, aux = jacobian(
+    jac = jacobian(
         fn,
         operator.in_size(),
         operator.out_size(),
         holomorphic=any(jnp.iscomplexobj(xi) for xi in jtu.tree_leaves(operator.x)),
-        has_aux=True,
         jac=operator.jac,
     )(operator.x)
-    out = PyTreeLinearOperator(jac, operator.out_structure(), operator.tags)
-    return AuxLinearOperator(out, aux)
+    return PyTreeLinearOperator(jac, operator.out_structure(), operator.tags)
 
 
 @materialise.register(FunctionLinearOperator)
@@ -1516,8 +1476,16 @@ def is_symmetric(operator: AbstractLinearOperator) -> bool:
 
 def _has_real_dtype(operator) -> bool:
     """Check if all dtypes in an operator's structure are real (not complex)."""
-    leaves = jtu.tree_leaves(operator.in_structure())
-    return all(jnp.issubdtype(leaf.dtype, jnp.floating) for leaf in leaves)
+    leaves = jtu.tree_leaves((operator.in_structure(), operator.out_structure()))
+    dtype = jnp.result_type(*leaves)
+    if jnp.issubdtype(dtype, jnp.complexfloating):
+        return False
+    elif jnp.issubdtype(dtype, jnp.floating):
+        return True
+    else:
+        assert False, (
+            "Only `jnp.floating` and `jnp.complexfloating` dtypes are understood."
+        )
 
 
 @is_symmetric.register(MatrixLinearOperator)
@@ -1874,10 +1842,6 @@ def _(operator, transform=transform):
     def _(operator, transform=transform):
         return transform(operator.operator) / operator.scalar
 
-    @transform.register(AuxLinearOperator)  # pyright: ignore
-    def _(operator, transform=transform):
-        return transform(operator.operator)
-
 
 @linearise.register(TangentLinearOperator)
 def _(operator):
@@ -1938,11 +1902,6 @@ def _(operator):
     return (diag / operator.scalar, lower / operator.scalar, upper / operator.scalar)
 
 
-@tridiagonal.register(AuxLinearOperator)
-def _(operator):
-    return tridiagonal(operator.operator)
-
-
 @linearise.register(ComposedLinearOperator)
 def _(operator):
     return linearise(operator.operator1) @ linearise(operator.operator2)
@@ -1983,10 +1942,6 @@ def _(operator):
     def _(operator, check=check):
         return check(operator.primal)
 
-    @check.register(AuxLinearOperator)
-    def _(operator, check=check):
-        return check(operator.operator)
-
 
 # Scaling/negating preserves these structural properties
 for check in (
@@ -2272,8 +2227,3 @@ def _(operator):
 @conj.register(ComposedLinearOperator)
 def _(operator):
     return conj(operator.operator1) @ conj(operator.operator2)
-
-
-@conj.register(AuxLinearOperator)
-def _(operator):
-    return conj(operator.operator)