Skip to content

Lazily import modules#1026

Open
kbattocchi wants to merge 3 commits into
mainfrom
kebatt/lazy-imports
Open

Lazily import modules#1026
kbattocchi wants to merge 3 commits into
mainfrom
kebatt/lazy-imports

Conversation

@kbattocchi

Copy link
Copy Markdown
Member

This PR adds a utility to lazily import modules, and uses it for the following use-cases:

  • Only load shap and statsmodels when needed; as opposed to critical imports like numpy or sklearn, they are only needed for a small number of features, but they add substantial time to importing our library
  • In places where we are trying to avoid circular dependencies between our own subpackages, this lets us bring the reference to the other package up to the top with other imports rather than putting it inside of the functions where the actual use occurs

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copilot encountered an error and was unable to review this pull request. You can try again by re-requesting a review.

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

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

Comment thread econml/utilities.py Outdated
kbattocchi and others added 3 commits June 11, 2026 17:01
Add a _LazyModule proxy class (econml/_lazy.py) that defers module
loading until first attribute access.  This keeps lazy import
declarations at the top of each file alongside normal imports,
making the deferred loading explicit and avoiding scattered inline
imports inside function bodies.

Modules deferred:
- shap (+numba, sparse) in econml/_shap.py
- statsmodels.iolib.{table,summary} in econml/utilities.py
- statsmodels.{tools,api,robust} in econml/sklearn_extensions/linear_model.py
- statsmodels.tools.tools in econml/data/dynamic_panel_dgp.py
- statsmodels.{api,tools} in econml/validate/drtester.py

Measured improvement: single-test cold start ~12s -> ~7s (39% faster).

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Keith Battocchi <kebatt@microsoft.com>
Add a lightweight add_constant() to econml/utilities.py that handles
array-like input (including numpy arrays and pandas DataFrame/Series)
directly. Unlike statsmodels' version it always returns an ndarray —
this is documented in the function's Notes section. This eliminates
the statsmodels dependency from dynamic_panel_dgp.py entirely, and
removes the _statsmodels_tools lazy import from linear_model.py and
drtester.py.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Keith Battocchi <kebatt@microsoft.com>
Replace 4 deferred imports that existed to avoid circular imports
with top-level _LazyModule declarations.  The lazy proxy defers the
actual importlib.import_module() call until first attribute access,
which happens inside function/method bodies after all modules have
finished loading — so the circular dependency is still broken, but
the import declaration lives at the top of the file.

- econml/dml/causal_forest.py: econml.score (RScorer)
- econml/inference/_bootstrap.py: econml._cate_estimator (BaseCateEstimator)
- econml/sklearn_extensions/linear_model.py: econml.sklearn_extensions.model_selection
- econml/_ortho_learner.py: econml.dml._rlearner (_ModelFinal)

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Keith Battocchi <kebatt@microsoft.com>
@kbattocchi kbattocchi force-pushed the kebatt/lazy-imports branch from baa6abb to d9719cb Compare June 11, 2026 21:34
@kbattocchi kbattocchi requested a review from fverac June 11, 2026 21:59
@kbattocchi kbattocchi marked this pull request as ready for review June 11, 2026 21:59
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants