How to use static args in while_loop, or to say how to use dynamic_slice in while_loop.

[Big-Baby]

When I prepared to use fixed points algorithms with jax, I found an open-source package fxp_jax, however, fxp_jax only implemented simple aitken_acceleration and fxp_jax do NOT use vmap. And I prefer to use anderson_acceleration.
For a simple LASSO regression, I implemented anderson_acceleration as the following:

from jax import lax
from jax import random
from jax import jit, vmap
from jax import numpy as jnp
from functools import partial

@partial(jit, static_argnames=["g", "m", "lam", "tol", "max_iter"])
def single_anderson_acceleration(g, x0, y, m, lam, tol, max_iter, beta = 1.0):
    x_shape = x0.shape
    x_entries = 1
    for v in x_shape:
        x_entries *= v
    X = jnp.zeros(shape=(m, x_entries))
    G = jnp.zeros(shape=(m, x_entries))
    X = X.at[0].set(x0.reshape((-1,)))
    X = X.at[1].set(g(x0, y).reshape((-1,)))
    G = G.at[0].set(X.at[1].get())
    G = G.at[1].set(g(X.at[1].get().reshape(x_shape), y).reshape((-1,)))

    H = jnp.ones(shape=(m + 1, m + 1))
    H = H.at[0, 0].set(0)
    H = H.at[1:, 1:].set(jnp.zeros(shape=(m, m)))

    z = jnp.zeros(shape=(m+1,1))
    z = z.at[0].set(1)

    res = jnp.full(max_iter, jnp.inf)

    @partial(jit, static_argnames=["i"])
    def _body_fn(i, x_hist, g_hist, h, residuals):
        t = (i + 2) % m
        n = i + 2
        g_slice = lax.dynamic_slice(
            g_hist,
            start_indices=(0, 0),
            slice_sizes=(n, x_entries)
        )
        x_slice = lax.dynamic_slice(
            x_hist,
            start_indices=(0, 0),
            slice_sizes=(n, x_entries)
        )
        u = g_slice - x_slice

        uut = u @ u.T + lam * jnp.eye(n)
        h = lax.dynamic_update_slice(
            h,
            uut,
            start_indices=(1, 1)
        )
        h_slice = lax.dynamic_slice(
            h,
            start_indices=(0, 0),
            slice_sizes=(n + 1, n + 1)
        )
        z_slice = lax.dynamic_slice(
            z,
            start_indices=(0,),
            slice_sizes=(n + 1,)
        )
        alpha = jnp.linalg.solve(h_slice, z_slice)[1:n + 1]
        term_1 = jnp.einsum("b,bn->n", alpha, g_hist.at[:n].get())
        term_2 = jnp.einsum("b,bn->n", alpha, x_hist.at[:n].get())
        x_t = beta * term_1 + (1. - beta) * term_2
        g_t = g(x_t, y).reshape((-1,))
        x_hist = x_hist.at[t].set(x_t)
        g_hist = g_hist.at[t].set(g_t)
        residual = jnp.linalg.norm(g_t - x_t)
        residuals = residuals.at[i].set(residual)
        return i + 1, x_hist, g_hist, h, residuals

    def body_fn(state):

        i, x_hist, g_hist, h, residuals = state
        return _body_fn(i, x_hist, g_hist, h, residuals)


    def cond_fn(state):
        i, x_hist, g_hist, h, residuals = state
        return jnp.logical_and(i < max_iter, residuals[i - 1] >= tol)

    iter_num, _, _, _, res = lax.while_loop(
        cond_fn,
        body_fn,
        (2, X, G, H, res)
    )
    return jnp.zeros_like(x0), iter_num, res

def anderson_acceleration(g, x0, y, m=5, lam=1e-4, tol=1e-6, max_iter=1000):
    x, max_iter, residuals = vmap(single_anderson_acceleration, in_axes=(None, 0, 0, None, None, None, None))(g, x0, y, m, lam, tol, max_iter)
    return x, max_iter, residuals.at[:max_iter].get()

def _test():
    key = random.key(42)
    a_key, key = random.split(key)
    A = random.normal(a_key, shape=(30, 60))
    A = A / jnp.sum(jnp.square(A), axis=0, keepdims=True)
    x_normal_key, key = random.split(key)
    x_normal = random.normal(x_normal_key, shape=(20, 60, 1))
    x_binomial_key, key = random.split(key)
    x_binomial = random.binomial(x_binomial_key, 2, 0.1, shape=(20, 60, 1))
    y = A @ (x_normal * x_binomial)
    alpha = 0.005
    lam = 0.1
    ss = jit(lambda x: jnp.sign(x) * jnp.maximum(jnp.abs(x) - alpha * lam, 0))
    g = jit(lambda x, y_: ss(x - alpha * A.transpose() @ (A @ x - y_)))
    x_pred, tot_iter, residuals = anderson_acceleration(g, jnp.zeros(shape=(20, 60, 1)), y, tol=1e-5)
    print(tot_iter, residuals)

if __name__ == '__main__':
    _test()

While, some errors occurred,:

Traceback (most recent call last):
  File "/home/user/PycharmProjects/PythonProject/algoritms/anderson.py", line 131, in <module>
    _test()
    ~~~~~^^
  File "/home/user/PycharmProjects/PythonProject/algoritms/anderson.py", line 127, in _test
    x_pred, tot_iter, residuals = anderson_acceleration(g, jnp.zeros(shape=(20, 60, 1)), y, tol=1e-5)
                                  ~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/sw/PycharmProjects/PythonProject/algoritms/anderson.py", line 110, in anderson_acceleration
    x, max_iter, residuals = vmap(single_anderson_acceleration, in_axes=(None, 0, 0, None, None, None, None))(g, x0, y, m, lam, tol, max_iter)
                             ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/PycharmProjects/PythonProject/algoritms/anderson.py", line 102, in single_anderson_acceleration
    iter_num, _, _, _, res = lax.while_loop(
                             ~~~~~~~~~~~~~~^
        cond_fn,
        ^^^^^^^^
        body_fn,
        ^^^^^^^^
        (2, X, G, H, res)
        ^^^^^^^^^^^^^^^^^
    )
    ^
  File "/home/user/PycharmProjects/PythonProject/algoritms/anderson.py", line 95, in body_fn
    return _body_fn(i, x_hist, g_hist, h, residuals)
ValueError: Non-hashable static arguments are not supported. An error occurred while trying to hash an object of type <class 'jax._src.interpreters.partial_eval.DynamicJaxprTracer'>, Traced<~int32[]>with<DynamicJaxprTrace>. The error was:
TypeError: unhashable type: 'DynamicJaxprTracer'

--------------------
For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.

[jakevdp]

The state within a while_loop is non-static by definition, so you will not be able to pass any of its components as a static argument. That is: i in _body_fun must be dynamic. It looks like this is not possible in your code, because you specify dynamic_slice output shapes based on the value of i. Unfortunately your code is not compatible with execution via while_loop for this reason, and if you want to proceed you'll need to rewrite it to only use statically-shaped computations.