Introduce JAX OLS Class

[s3alfisc]

• Move all jax code to a dedicated JAX module.
• Add a basic JAXOLS class.

from pyfixest.estimation.jax.OLSJAX import OLSJAX
import jax.numpy as jnp
from jax.random import PRNGKey
import jax

N = 1000
k = 3
X = jax.random.normal(PRNGKey(0), shape=(N, k))
beta = jnp.array([1, 2, 3])
Y = (X @ beta + jax.random.normal(PRNGKey(0), shape=(N,))).reshape(-1, 1)
fe = jax.random.randint(PRNGKey(0), minval=0, maxval=5, shape=(N,1))

ols = OLSJAX(Y = Y, X = X, fe = fe, vcov = "HC1")
ols.fit()
ols.tidy()

Idea: implement all computational steps in pure JAX. All non-directly implemented methods can be inherited from Feols (i.e. tidy, inference, ritest, etc unless directly implemented).

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[juanitorduz]

docstrings?

[juanitorduz]

LGTM up to the docstrings ;)

[apoorvalal]

this file should probably be deleted in this PR

[apoorvalal]

this probably supercedes the numpy API I requested a while ago, provided we keep the methods+signatures consistent with statsmodels on the base OLSJAX class and wrap them appropriately through pf.feols.

[s3alfisc]

provided we keep the methods+signatures consistent with statsmodels on the base OLSJAX class and wrap them appropriately through pf.feols.

Yes, that was exactly the idea! I think some minor refactoring of the different OLS classes will be required (for example, we might have to move all the method calls from FixestMulti into the get_fit() methods, though that's something I wanted to do anyways for the numpy API). Then I was thinking about adding a Interface class that inherits that will be called from FixestMulti but inherits from Feols and whose get_fit method is overwritten by the JAX method. Might be a bad design though? It for sure sounds more complex than it will be, I'll push a PR in the next days, then it will be easier to see what I'd like to do.

[apoorvalal]

curious whether using lx.linear_solve gets speed-gains here; example

[s3alfisc]

Speed gains look superb, how do you feel about adding lineax as an optional dependency to the JAX env?

[apoorvalal]

i'm game; Kidger's stuff is very high quality and stable, so I doubt that it will break anytime soon.

[apoorvalal]

i'll test out lineax's solver on this PR over the weekend; expect a commit

[apoorvalal]

if residuals and scores methods are simply setting the properties on this class, do they need return statements? alternatively, if they do return arrays, set them inside fit as self.uhat = self.residuals ?

[s3alfisc]

Yes, I agree with you, I think simply doing something as

    @property
    def residuals(self):
        return self.Y - self.X @ self.beta

and then potentially setting an uhat attribute in fit() would be cleaner? Optimally we can completely drop the uhat attribute, I don't think it is needed per se (except for some compatibility with the current implementation).

[apoorvalal]

makes sense. re: not setting it as an attribute - agree; i err on the side of not attaching too many things to the model class because it ends up being a memory-consumption headache. Given that the benefits of jax are primarily for big data ™️, making the model class as lean as possible by not attaching any data (that's already built into pf.feols right) would make sense to me

[s3alfisc]

2c92c3d implements a POC integration of the OLSJAX class into the pyfixest stream. There are still plenty of things to polish and implement (vcov options, feols API, multicollinearity checks etc). But it should provide you with an idea how we could generally integrate "proper numpy APIs" into the current pyfixest internal design (irrespective of JAX integration).

I am not sure if what I put together is a masterclass in software engineering, so I'd love to hear your opinions @apoorvalal @juanitorduz . Moving all other estimation classes to such a design would be a bit of work - as a benefit, we'd get a portable estimation classes with numpy / JAX APIs. I still have to ponder about this a bit.

For now, I will open a new PR that simply creates the JAX module to unblock @iamlemec's PR to add his performance-improved implementation of the demeaning algo, as merging the OLSJAX class and an appropriate interface will require a little bit more work & thinking.

[juanitorduz]

Personally, I think it's better to aim for "good" instead of "perfect" and work on smaller iterations. So, I do not have any objections to merging an initial version and creating issues to work on iterations. Just food for though ;) (we just need to be conscious of the technical debt, but I think at his earlier stage the risk is low)

[s3alfisc]

Personally, I think it's better to aim for "good" instead of "perfect" and work on smaller iterations. So, I do not have any objections to merging an initial version and creating issues to work on iterations

Oh I'm glad to hear! For me there's not much worse than a never ending PR that eventually needs multiple rebases to align with the main branch. Horror! 😄

So I'll proceed step by step:

• finalize the OLS class
• write a wrapper class and make OLSJAX callable via feols() / FixestMulti
• port the multicollinearity checks to JAX

After this is set up, we can consider if the wrapper class actually looks decent, or if it will be hard to maintain in the long term and iterate?