How to use the causallib.utils.stat_utils.robust_lookup function in causallib

Will clip probabilities between clip_eps and 1-clip_eps.
            use_stabilized (None|bool): Whether to re-weigh the learned weights with the prevalence of the treatment.
                                        This overrides the use_stabilized parameter provided at initialization.
                                        If True provided, but the model was initialized with use_stabilized=False, then
                                        prevalence is calculated from data at hand, rather than the prevalence from the
                                        training data.
                                        See Also: https://www.ncbi.nlm.nih.gov/pmc/articles/PMC4351790/#S6title

        Returns:
            pd.Series | pd.DataFrame: If treatment_values is not supplied (None) or is a scalar, then a vector of
                                      n_samples with a weight for each sample is returned.
                                      If treatment_values is a list/array, then a DataFrame is returned.
        """
        weight_matrix = self.compute_weight_matrix(X, a, truncate_eps, use_stabilized)
        if treatment_values is None:
            weights = robust_lookup(weight_matrix, a)  # lookup table: take the column a[i] for every i in index(a).
        else:
            weights = weight_matrix[treatment_values]
        return weights

def _estimator_predict(self, X, a):
        """Predict on data.

        Args:
            X (pd.DataFrame): Covariates.
            a (pd.Series): Target variable - treatment assignment

        Returns:
            PropensityEvaluatorPredictions
        """
        propensity = self.estimator.compute_propensity(X, a, treatment_values=a.max())
        propensity_by_treatment_assignment = self.estimator.compute_propensity_matrix(X)
        propensity_by_treatment_assignment = robust_lookup(propensity_by_treatment_assignment, a)

        weight_prediction = super(PropensityEvaluator, self)._estimator_predict(X, a)
        # Do not force stabilize=False as in WeightEvaluator:
        weight_by_treatment_assignment = self.estimator.compute_weights(X, a)
        prediction = PropensityEvaluatorPredictions(weight_by_treatment_assignment,
                                                    weight_prediction.weight_for_being_treated,
                                                    weight_prediction.treatment_assignment_pred,
                                                    propensity,
                                                    propensity_by_treatment_assignment)
        return prediction