Feature Request: SQUAREM Accelerated Fixed Point Methods

Thanks for the fantastic library! I have a memory sensitive application that requires me to prefer iterative methods for my fixed point solve rather than Jacobian based methods. SQUAREM is a fixed point acceleration method that effectively approximates the Jacobian sans memory requirements.

Before the advent of optimistix I effectively implemented SQUAREM with an optimistix-like interface, sans the actually robust software engineering (I am an economist with minimal formal software engineering training). I'm looking to upstream SQUAREM within optimistix as a replacement.

Description of Algorithm

SQUAREM modifies the naive iteration update step with the following. For a given fixed point equation $F(\theta)$, run the function $F$ twice. Then apply a steplength $\alpha$ based on these two updates, and run the function a final time. The pseudocode references the original 2008 SQUAREM paper, which unfortunately is not open access.

The authors suggest three possibilities for the steplength $\alpha$, S1 to S3. S3 is recommended.

Proposed Implementation

I've implemented this by modifying the _solver/fixed_point.py file to:

1. Include S1-S3 within the class
2. Implement the pseudocode above for the step function

This is my first foray into playing with the internals of a "proper" library and

1. I really have no clue what the $\omega$ in the code are supposed to do other than it's imported from equinox. I added them everywhere until the code stopped complaining.
2. I wasn't sure how to initialise the stepsize as a parameter (which it should be), so like any amateur code conjurer I've hard-coded the update function to use scheme 3.

class Squarem(
    AbstractFixedPointSolver[Y, Aux, _FixedPointState], strict=True
):
    """
    Uses SQUAREM acceleration
    """
    # define the different schemes for SQUAREM updating
    def s1(r: jnp.ndarray, v: jnp.ndarray):
        r"""Scheme 1 of SQUAREM algorithm alpha

        This is scheme S1 (equation (7) of Varadham and Roland (2008)) for the computation of alpha. :math:`r^Tv/v^Tv`. Notation below for parameters is as per the original paper.

        Parameters
        ----------
        r : jnp.ndarray
            :math:`\theta_1 - \theta_0`
        v : jnp.ndarray
            :math:`(\theta_2 - \theta_1) - r`
        """
        return (ω(r).call(jnp.transpose) @ v**ω).ω / (ω(v).call(jnp.transpose) @ v**ω).ω
    
    def s2(r: jnp.ndarray, v: jnp.ndarray):
        r"""Scheme 2 of SQUAREM algorithm alpha

        This is scheme S2 (equation (8) of Varadham and Roland (2008)) for the computation of alpha. :math:`r^Tr/r^Tv`. Notation below for parameters is as per the original paper.

        Parameters
        ----------
        r : jnp.ndarray
            :math:`\theta_1 - \theta_0`
        v : jnp.ndarray
            :math:`(\theta_2 - \theta_1) - r`
        """
        return (ω(r).call(jnp.transpose) @ r**ω).ω / (ω(r).call(jnp.transpose) @ v**ω).ω
    
    def s3(r: jnp.ndarray, v: jnp.ndarray):
        r"""Scheme 3 of SQUAREM algorithm alpha

        This is scheme S3 (equation (9) of Varadham and Roland (2008)) for the computation of alpha. :math:`r^Tr/r^Tv`. Notation below for parameters is as per the original paper. This is the recommended scheme.

        Parameters
        ----------
        r : jnp.ndarray
            :math:`\theta_1 - \theta_0`
        v : jnp.ndarray
            :math:`(\theta_2 - \theta_1) - r`
        """
        return -1 * (ω(r).call(lambda x: jnp.linalg.norm(x, ord = 2)) / ω(v).call(lambda x: jnp.linalg.norm(x, ord = 2))).ω
    
    rtol: float
    atol: float
    norm: Callable[[PyTree], Scalar] = max_norm

    def init(
        self,
        fn: Fn[Y, Y, Aux],
        y: Y,
        args: PyTree,
        options: dict[str, Any],
        f_struct: PyTree[jax.ShapeDtypeStruct],
        aux_struct: PyTree[jax.ShapeDtypeStruct],
        tags: frozenset[object],
    ) -> _FixedPointState:
        del fn, y, args, options, f_struct, aux_struct
        return _FixedPointState(jnp.array(jnp.inf))

    def step(
        self,
        fn: Fn[Y, Y, Aux],
        y: Y,
        args: PyTree,
        options: dict[str, Any],
        state: _FixedPointState,
        tags: frozenset[object],
        scheme = s3,
    ) -> tuple[Y, _FixedPointState, Aux]:
        # step twice
        y1, _ = fn(y, args)
        y2, aux = fn(y, args)
        # compute r, v in the paper
        r = (y1**ω - y**ω).ω
        v = (y2**ω - y**ω).ω
        # compute the alpha based on scheme
        alpha = scheme(r, v)
        y_prime = (y**ω - 2 * alpha**ω * r**ω + ω(alpha).call(lambda x: jax.lax.integer_pow(x, 2)) * v**ω).ω
        # last stabilisation step
        y_new, aux = fn(y_prime, args)
        error = (y**ω - y_new**ω).ω
        with jax.numpy_dtype_promotion("standard"):
            scale = (self.atol + self.rtol * ω(y_prime).call(jnp.abs)).ω
            new_state = _FixedPointState(self.norm((error**ω / scale**ω).ω))
        return y_new, new_state, aux

If there's interest, I'm happy to do up a PR, but I'll need quite a bit of guidance along the way.

Start with _solver/fixed_point.py and compare the proposed Squarem implementation with the existing fixed-point solver interface. Resolve how Equinox ω wrappers, the stepsize scheme, state, and auxiliary values should be handled. Done means SQUAREM is integrated with a sound interface for selecting S1–S3 and implementing the described three-evaluation update.