package this into a usable library

[apoorvalal]

linear, GMM, and causal submodules.

Regression demo; run on beefier machine for the ill posed case.

[Copilot]

Pull Request Overview

This PR packages the core econometrics routines into a standalone jaxonometrics library, adds tests for the new linear and GMM modules, and provides configuration and demonstration notebooks.

• Introduce LinearRegression, GMM (GMM, LinearIVGMM, TwoStepGMM), and causal (EntropyBalancing) estimators
• Add unit tests for linear regression and IV-GMM
• Supply a pyproject.toml for package metadata and multiple demo notebooks

Reviewed Changes

Copilot reviewed 23 out of 28 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
tests/test_linear.py	Add test comparing LinearRegression vs. scikit-learn
tests/test_gmm.py	Add tests for GMM and LinearIVGMM estimators
pyproject.toml	Define package metadata and dependencies
jaxonometrics/linear.py	Implement LinearRegression with Lineax solver
jaxonometrics/gmm/base.py	General-purpose GMM class with LM and SGD support
jaxonometrics/gmm/linear_iv.py	LinearIVGMM convenience wrapper
jaxonometrics/gmm/twostep.py	TwoStepGMM for nuisance-parameter models
jaxonometrics/causal.py	EntropyBalancing causal reweighting estimator
jaxonometrics/base.py	Abstract BaseEstimator class
jaxonometrics/legacy/**	Legacy notebooks and scripts (benchmarks and demos)

Comments suppressed due to low confidence (4)

jaxonometrics/gmm/twostep.py:37

• The signature uses init_theta: Dict and init_eta: Dict but downstream LevenbergMarquardt expects numeric arrays. Consider annotating and requiring jnp.ndarray for both initial parameter arrays.

    def fit(

jaxonometrics/causal.py:1

• No unit tests target the EntropyBalancing estimator. Add tests to verify that fit computes correct weights on a simple synthetic dataset.

from typing import Dict, Optional

jaxonometrics/gmm/base.py:62

• The OptaxSolver.update method returns the updated parameters and state, not the raw gradients. Rename grads to new_params and assign params = new_params directly instead of calling optax.apply_updates.

                grads, opt_state = solver.update(params=params, state=opt_state, data=data)

pyproject.toml:16

• [nitpick] Dependency narwhals does not appear to be used in the code. Remove or document its purpose to avoid unnecessary bloat.

  "narwhals",

[Copilot]

[nitpick] The legacy folder includes many notebooks and scripts that clutter the package. Consider moving these to a separate examples or benchmarks repository to keep the core library lean.