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

def imputer_test(self):

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

        extract_features = ['a']
        feature_extractor = FeatureExtractor(input_scalars=['a'],
                                         output_vector='extracted_a_output',
                                         output_vector_items=["{}_out".format(x) for x in extract_features])

        imputer = Imputer(strategy='mean')
        imputer.mlinit(prior_tf=feature_extractor,
                       output_features='a_imputed')

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

        imputer.fit(df2[['a']])

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

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

        expected_model = {
          "op": "imputer",
          "attributes": {