EnsmallenEstimator

EnsmallenEstimator(moment_cond, weighting_matrix='optimal')

Generalized method of moments estimator with JAX-powered derivatives.

The estimator uses JAX to evaluate gradients and Jacobians while relying on the pyensmallen optimization backend for numerical minimization. It supports two-step GMM through an optimal weighting matrix update and provides asymptotic as well as bootstrap-based uncertainty estimates.

Methods

Name	Description
bootstrap_full	Compute bootstrap standard errors using batched processing.
bootstrap_scores	Compute bootstrap standard errors using the fast score bootstrap method.
compute_asymptotic_variance	Compute the asymptotic covariance matrix and standard errors.
fit	Fit the GMM model using PyEnsmallen optimizer with JAX gradients.
gmm_objective	Compute GMM objective function value and gradient for PyEnsmallen.
iv_moment	Standard IV moment function: z_i * (y_i - x_i’β).
jacobian_moment_cond	Compute Jacobian of moment conditions using JAX.
optimal_weighting_matrix	Compute the optimal GMM weighting matrix.
summary	Generate summary statistics for the fitted model.

bootstrap_full

EnsmallenEstimator.bootstrap_full(
    n_bootstrap=1000,
    seed=None,
    batch_size=50,
    verbose=False,
)

Compute bootstrap standard errors using batched processing.

Parameters

Name	Type	Description	Default
n_bootstrap	int	Number of bootstrap draws.	1000
seed	int or None	Random seed for reproducibility.	None
batch_size	int	Number of bootstrap problems processed together.	50
verbose	bool	Whether to print bootstrap progress.	False

Returns

Name	Type	Description
	ndarray of shape (n_params,)	Bootstrap standard errors for the fitted parameter vector.

bootstrap_scores

EnsmallenEstimator.bootstrap_scores(n_bootstrap=1000, seed=None, verbose=False)

Compute bootstrap standard errors using the fast score bootstrap method.

This method is computationally more efficient than the full bootstrap. It bootstraps the score functions (moment conditions) rather than recomputing the GMM estimator for each bootstrap sample.

Parameters

Name	Type	Description	Default
n_bootstrap	int	Number of bootstrap draws.	1000
seed	int or None	Random seed for reproducibility.	None
verbose	bool	Whether to print bootstrap progress.	False

Returns

Name	Type	Description
	ndarray of shape (n_params,)	Bootstrap standard errors for the fitted parameter vector.

compute_asymptotic_variance

EnsmallenEstimator.compute_asymptotic_variance()

Compute the asymptotic covariance matrix and standard errors.

Returns

Name	Type	Description
	None	Updates vtheta_ and std_errors_ in place.

fit

EnsmallenEstimator.fit(z, y, x, verbose=False)

Fit the GMM model using PyEnsmallen optimizer with JAX gradients.

Parameters

Name	Type	Description	Default
z	ndarray of shape (n_samples, n_instruments)	Instrument matrix.	required
y	ndarray of shape (n_samples,)	Outcome vector.	required
x	ndarray of shape (n_samples, n_params)	Regressor matrix. Include a column of ones if an intercept is desired.	required
verbose	bool	Whether to print optimization and inference warnings.	False

Returns

Name	Type	Description
	None	Stores fitted parameters and inference results on the estimator.

gmm_objective

EnsmallenEstimator.gmm_objective(beta, gradient)

Compute GMM objective function value and gradient for PyEnsmallen.

Parameters

Name	Type	Description	Default
beta	ndarray of shape (n_params,)	Candidate parameter vector.	required
gradient	ndarray of shape (n_params,)	Output array filled in-place with the gradient of the criterion.	required

Returns

Name	Type	Description
	float	Value of the GMM objective at beta.

iv_moment

EnsmallenEstimator.iv_moment(z, y, x, beta)

Standard IV moment function: z_i * (y_i - x_i’β).

Works with both NumPy and JAX arrays.

Parameters

Name	Type	Description	Default
z	ndarray	Instrument matrix.	required
y	ndarray	Outcome vector.	required
x	ndarray	Regressor matrix.	required
beta	ndarray	Parameter vector.	required

Returns

Name	Type	Description
	ndarray	Matrix of per-observation moment conditions.

jacobian_moment_cond

EnsmallenEstimator.jacobian_moment_cond()

Compute Jacobian of moment conditions using JAX.

Returns

Name	Type	Description
	ndarray of shape (n_moments, n_params)	Jacobian of the mean moment vector with respect to the parameters.

optimal_weighting_matrix

EnsmallenEstimator.optimal_weighting_matrix(moments, epsi=1e-08)

Compute the optimal GMM weighting matrix.

Parameters

Name	Type	Description	Default
moments	ndarray of shape (n_samples, n_moments)	Per-observation moment evaluations at the current parameter value.	required
epsi	float	Diagonal regularization added before inversion for numerical stability.	1e-8

Returns

Name	Type	Description
	ndarray of shape (n_moments, n_moments)	Estimated optimal weighting matrix.

summary

EnsmallenEstimator.summary(prec=4, param_names=None)

Generate summary statistics for the fitted model.

Parameters

Name	Type	Description	Default
prec	int	Number of decimal places used in the reported table.	4
param_names	list of str or None	Optional parameter labels. When omitted, parameters are labeled theta_0, theta_1, and so on.	None

Returns

Name	Type	Description
	DataFrame	Summary table containing coefficients, standard errors, test statistics, p-values, and confidence intervals.