How to use the mleap.sklearn.extensions.data.Imputer function in mleap

def imputer_ext_test(self):

        def _set_nulls(df):
            row = df['index']
            if row in [2,5]:
                return np.NaN
            return df.a

        imputer = Imputer(strategy='mean', input_features='a', output_features='a_imputed')

        df2 = self.df
        df2.reset_index(inplace=True)
        df2['a'] = df2.apply(_set_nulls, axis=1)

        imputer.fit(df2)

        self.assertAlmostEqual(imputer.statistics_[0], df2.a.mean(), places = 7)

        imputer.serialize_to_bundle(self.tmp_dir, imputer.name)

        expected_model = {
            "op": "imputer",
            "attributes": {
                "surrogate_value": {
                    "double": df2.a.mean()