How to add hierarchical/multi-level model aspects to Multi-Investment Attribution Case Study

All of the use cases of causal modeling I'm involved include a geographic component in which it's critical to be able to simultaneously estimate both parent-level insights (like Total US) and child-level insights (like States). If the Multi-Investment Attribution Case Study were modified to include 2 geographic indicator columns (1 for zip code , 1 for state), and you wanted to get both global and state-level causal results, how would you do it in a way that applies partial pooling (like mixed modeling with Pymer4 does) and not just a For-loop by state?

Here is the the original Case Study excerpt. Can this formulation (or the data) be modified to allow mixed modeling?

Train EconML model with generic helper models

model = LinearDRLearner(
model_regression=XGBRegressor(learning_rate=0.1, max_depth=3),
model_propensity=XGBClassifier(learning_rate=0.1, max_depth=3, objective="multi:softprob"),
random_state=1,
)

Specify final stage inference type and fit model

model.fit(Y=Y, T=T, X=X, W=W, inference="statsmodels")

Begin with the Multi-Investment Attribution Case Study and its LinearDRLearner example; review how X, W, Y, T, and inference are supplied to fit. Determine whether parent- and child-level geographic partial pooling is supported or requires a design proposal, with completion defined as a documented, reproducible approach rather than a per-state loop.