fepois not honoring demean_backend?

[christian-wilson]

I had been dealing with some issues where Numba failed while running fepois, so I thought using jax might solve the issue. However, when switching the demean backend, nothing seemed to change.

It seems like fepois uses the demean function from pyfixest.estimation.demean_, which only uses numba. I believe that in order for jax to be used, this should be demean_model?

New here, and my first issue, so could very much be wrong on this. Just basing this off of my experience and looking at this code in fepois_.py:

from `pyfixest.estimation.demean_ import demean

and inside the get_fit function

if _fe is not None:
    # ZX_resid = algorithm.residualize(ZX, mu)
    ZX_resid, success = demean(
        x=ZX, flist=_fe, weights=mu.flatten(), tol=_fixef_tol
    )

[s3alfisc]

Hi @crwilson87 , that's right, fepois is not yet supporting the JAX demeaner backend. But that's very easy to change - basically here, you'd have to change calling demean to JAX, e.g. by defining a demean_func. And then of course pass function arguments from from the fepois API to the FixestMulti to the Fepois class.

Would you feel up to try this? Else I can also add support in the next days.

Can you tell me a little more about the numba errors you are experiencing?

Also, not that if you're not running on the GPU, the JAX demeaner will likely be quite a bit slower than numba.