Skip to content

Lazily import modules#1026

Draft
kbattocchi wants to merge 3 commits intomainfrom
kebatt/lazy-imports
Draft

Lazily import modules#1026
kbattocchi wants to merge 3 commits intomainfrom
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

kbattocchi and others added 3 commits April 10, 2026 13:43
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
the numpy-array case directly, with a guard that raises TypeError
for pandas DataFrames.  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>
Copy link
Copy Markdown

Copilot AI left a comment

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.

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.

2 participants