sklearn SimpleImputer using mean from different column groups?

I’m looking at the SimpleImputer, in particular in here, and I would like to do the imputation on different columns. My data has 3 different sample groups, and I would like to do mean imputation by sample group rather than column-wise as done with SimpleImputer. My end goal is to then compute MSE through crossvalidation in the same way it is done here. Is this possible?

I could impute values using standard dataframe manipulation, but don’t know how to then use it for crossvalidation.

Thanks in advance