Sparse matrix features for DMLIV

I have been trying to use DMLIV for a problem where we want to estimate fixed effects of many features, in our case about 75000 of them. We have about 5.5 million observations. I have to use sparse matrices for the features to avoid running out of memory (256GB). What is the proper way to use DMLIV with sparse matrix features?

To illustrate what I have tried, the data looks something like

n = 5493141
p = 75000
y = np.random.poisson(size=(n, 1))
t = np.random.binomial(1, .001, size=(n, 1))
z = np.random.binomial(1, .001, size=(n, 1))
z_spmat = scipy.sparse.csr_matrix(z)

density = 0.00000001
size = int(n * p * density)

rows = np.random.randint(0, 2, size=size)
cols = np.random.randint(0, 2, size=size)
data = np.random.rand(size)

x_spmat = scipy.sparse.csr_matrix((data, (rows, cols)), shape=(n, p))
x_df = pd.DataFrame.sparse.from_spmatrix(x)

I am using SGD classifiers/regressors for first stage models, which I am able to fit to my data outside of econml:

from sklearn.linear_model import SGDClassifier, SGDRegressor

occ_sgd3 = DMLIV(
  model_Y_X = SGDRegressor(),
  model_T_X = SGDClassifier(loss='log'),
  model_T_XZ = SGDClassifier(loss='log'),
  model_final = SGDRegressor(),
)

When I pass in the features x as a sparse matrix:

occ_sgd3.fit(Y=y, T=t, Z=z, X=x_spmat)

I get the error IndexError: tuple index out of range.

When I pass them in as a pandas data frame with sparse columns

occ_sgd3.fit(Y=y, T=t, Z=z, X=x_df)

I get ValueError: all the input arrays must have same number of dimensions, but the array at index 0 has 2 dimension(s) and the array at index 1 has 1 dimension(s).

After digging around a little I tried making z a sparse matrix, since it seems to get concatenated with x when fitting, and trying

occ_sgd3.fit(Y=y, T=t, Z=z_spmat, X=x_df)

gives ValueError: Found input variables with inconsistent numbers of samples: [2, 5493141].

Finally I tried to avoid cross fitting by setting cv=1; I seemed to get further and actually estimated the first stage models, but eventually get the error TypeError: 'coo_matrix' object is not subscriptable. It looks like there's code in utilities.py that handles sparse matrix input so I suspect my input formats are not quite right.

Start with the DMLIV.fit entry point and the sparse-input handling in utilities.py, then reproduce the reported fits with x_spmat, x_df, z_spmat, and cv=1. Trace the shape and indexing failures; done means establishing a supported sparse input format or a clearly scoped correction for these errors.