Custom pooling and Autodiff

[DiagRisker]

Hello Again,
I am trying to force custom pooling procedure instead of max, min, avg

@jax.jit
def Mixpmeanpool(*a,la=.5): # a dummy pooling proposition
    a = jnp.array(a)
    return la*jnp.log(jnp.exp(a).sum()) 

# pooling wrapper very close to linen (and nnx)
@Partial( jax.jit, static_argnames = ['reducer','Spdim', 'scale',])
def pooler( x, Spdim = 1, scale = 2, reducer = Partial(Mixpmeanpool, la =.4) ):
    strides = (1,) + (scale,)*Spdim + (1,)*(x.ndim-Spdim-1)
    dargs = dict(window_strides = strides, window_dimensions = strides, padding = ((0,0),)*x.ndim)
    return lax.reduce_window(x , init_value = 0., computation = reducer, **dargs)

# gradient wrapper
def Gjacob(f ):
    """ General Gradient of respective parameters (projecting jacobian on canonical directions (tangent valued of ones)  
        Usage : 
            direct gradient:
                Gjacob(fun , *args, wrapper = False) (input)
    """
    def Gacob(*x):
        y, vjp_fn = jax.vjp(f, *x)
        return vjp_fn(jnp.ones_like(y)) # tangent values at 1
    return Gacob


input = jax.random.uniform( key = rng(), shape = (1,6,1) )
YA = pooler(input, Spdim = 1) ; print("YA ", YA.shape)
Gtest = Gjacob( Partial(Mixpmeanpool, la=.4))(input) # the pooling proposition is differentiable
Gtest = Gjacob(Partial(pooler, Spdim=1, scale=2, reducer = Partial(Mixpmeanpool, la =.4)))(input) # something is missing

Note that feeding a pooling operator to lax.reduce_window, will lead to at least 2 input at a time inside the function Mixpmeanpool, hence the *a declaration!
from what I understand, theses are rows of data with respect to the window .

Does someone know how to reexpress this without breaking the autodiff chain ?
for classical nnx/nn . avg_pool / min_pool / max_pool , we have https://flax.readthedocs.io/en/v0.5.3/_modules/flax/linen/pooling.html, with lax.add, lax.max , lax.min