SubsetWrapper and DMLIV

@kbattocchi @vsyrgkanis I open a new issue on that to make less confusion. So, I am creating X for performing cate on DMLIV as an array and trying to refit the final model with a subset of that X. To do that I employed SubsetWrapper which takes as argument a feature_inds coming from a pandas database?
My X is XW (dimension (517,76)) which is defined as
XW = hstack([X[:, :(n_products)], W])
What I would like to do is imply refit the final model on just X[:, :(n_products)] (which has dim (517,11)) which I am doing by defining XW as a database and taking the first n_products columns from it. The point is that I do not understand what SubsetWrapper is going to take when I specify the feature_indices. Yet my original X term is an array, XW which I translated into pandas just for sake of simplicity. But I don't know if SubsetWrapper is going to take the subset of values of XW or what else...I am really confused about it sorry.

Is that correct?
I report here the dimensions of my cate.fit() arguments which according to the help are compatible:

print(Y.shape, T[:,0].shape, XW.shape, Z.shape)

(517,) (517,) (517, 76) (517, )

I am specifically referring to the example reported in NLSYM_Linear. There DMLIV is applied on bumpy arrays having dimensions:

print(y.shape, T.shape, X.shape, Z.shape)
(2991,) (2991,) (2991, 22) (2991,)

and Subsetwrapper has been applied on DRIV as follows:

# We could also fit a projection on a subset of the features by using the
# subset wrapper from our utilities.

# Example: including everything for expository purposes, but any array-like of indices would work
subset_names = set(['motheduc'])
# list of indices of features X to use in the final model
feature_inds = np.argwhere([(x in subset_names) for x in X_df.columns.values]).flatten()
print(feature_inds)
# Because we are projecting to a low dimensional model space, we can
# do valid inference and we can use statsmodel linear regression to get all
# the hypothesis testing capability
proj_driv_model_effect = lambda: SubsetWrapper(StatsModelLinearRegression(),
                                          feature_inds # list of indices of features X to use in the final model
                                         )

which I basically reproduced except that refit_final now does not support proj_driv_model_effect() as argument but only inference='auto'.

What is actually puzzling me is that refitting the final raises the following error:

matmul: Input operand 1 has a mismatch in its core dimension 0, with gufunc signature (n?,k),(k,m?)->(n?,m?) (size 11 is different from 78)

which does not make sense to me since the error basically tells that there is some problem in some multiplication. But actually the displayed sizes are the ones of the subset of features I would like to take (11) and the size of the original features increased by 2 (76+2). To provide further info I could say that the nuisances parameters re-used in refit_final are 2 (why not 3?) and have both size:

cate._cached_values.nuisances[0].shape
(517,)

as expected.

Start with the NLSYM_Linear example and trace SubsetWrapper, refit_final, and cate.fit using the reported shapes and feature_inds. Reproduce the 11-versus-78 matmul error, then determine whether the wrapper and refit path agree on the selected features and nuisance values; done means the behavior is explained and the reported refit case is resolved or documented.