`solver = jax` leads to errors with CRV standard errors

[s3alfisc]

Example:

import pyfixest as pf 
data = pf.get_data()

# no error
pf.feols("Y ~ X1 | f1 + f2", data=data, solver = "jax", vcov = "iid")
# no error
pf.feols("Y ~ X1 | f1 + f2", data=data, solver = "jax", vcov = "hetero")
# error 
pf.feols("Y ~ X1 | f1 + f2", data=data, solver = "jax", vcov = {"CRV1":"f1"})

# File "pyfixest\estimation\vcov_utils.py", line 131:
# def _crv1_meat_loop(
#     <source elided>
# ) -> np.ndarray:
#    k = scores.shape[1]

I suppose that there is some incompatibility between the numba code in the _crv1_meat_loop and jax?

[shapiromh]

I also ran into this while playing with the jax solver. I am not sure what other errors would come up further along, but at least this one seems to be caused by _scores being read by numba as some type of object that would not have the shape attribute (even though _scores does in reality have this attribute). Specifically, the error is "Unknown attribute 'shape' of type readonly buffer(float32, 2d, C)" (a known incompatibility).

Running the code in the linked thread seems to fix the first issue, but then it will fail on dtype (second line of the _crv1_meat_loop) function. There may be a way to fix this too using numba's attribute overloading (I tried ... but am not familiar enough with numba's underlying workings) but seems like this should be beyond the scope of a downstream module.

Unfortunately, it doesn't seem like that issue I shared has much attention so better integration coming from numba might be unlikely. It may be better to limit where jax arrays are passed around (not sure what the performance cost would be) or have separate implementations (I would guess there is a "jax way" to write the same numba functions) for the same interface.

[s3alfisc]

Hi @shapiromh , thanks so much for taking a look and actually digging out the JAX issue! Knowing about it will surely save me a lot of time going forward as I won't even try to fix this within PyFixest. Indeed it looks like the issue won't solved and I also agree that having a full "Jax flow" seems like the right approach. I've mapped out a very POC JaxOLS class in this PR #790 It's one of the things I'd love to get back to after releasing the next version; please let me know in case you want to tinker around with it a bit! Best, Alex

[shapiromh]

I'm a heavier and heavier user so I'd probably search around for where I can help add functionality after your next release, but JAX might be a fun way to contribute too.

[s3alfisc]

I love to hear it! Looking forward to any potential contributions =)