# Buggy code
def _count(self, X, Y):
    def _update_cat_count_dims(cat_count, highest_feature):
        diff = highest_feature + 1 - cat_count.shape[1]
        if diff > 0:
            # we append a column full of zeros for each new category
            return np.pad(cat_count, [(0, 0), (0, diff)], "constant")
        return cat_count

    def _update_cat_count(X_feature, Y, cat_count, n_classes):
        for j in range(n_classes):
            mask = Y[:, j].astype(bool)
            if Y.dtype.type == np.int64:
                weights = None
            else:
                weights = Y[mask, j]
            counts = np.bincount(X_feature[mask], weights=weights)
            indices = np.nonzero(counts)[0]
            cat_count[j, indices] += counts[indices]

    self.class_count_ += Y.sum(axis=0)
    for i in range(self.n_features_):
        X_feature = X[:, i]
        self.category_count_[i] = _update_cat_count_dims(
            self.category_count_[i], X_feature.max()
        )
        _update_cat_count(
            X_feature, Y, self.category_count_[i], self.class_count_.shape[0]
        )


# Traceback (IndexError)
---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
<ipython-input-3-551bb8080923> in <module>
10 model = CategoricalNB().fit(X_train, y_train)
11
---> 12 model.predict(X_test)

~/Documents/MachineLearning/onnx_projects/skl_env/lib/python3.6/site-packages/sklearn/naive_bayes.py in predict(self, X)
75         check_is_fitted(self)
76         X = self._check_X(X)
---> 77         jll = self._joint_log_likelihood(X)
78         return self.classes_[np.argmax(jll, axis=1)]
79

~/Documents/MachineLearning/onnx_projects/skl_env/lib/python3.6/site-packages/sklearn/naive_bayes.py in _joint_log_likelihood(self, X)
1217         for i in range(self.n_features_):
1218             indices = X[:, i]
-> 1219             jll += self.feature_log_prob_[i][:, indices].T
1220         total_ll = jll + self.class_log_prior_
1221         return total_ll

IndexError: index 5 is out of bounds for axis 1 with size 5

# Fixed code
def _count(self, X, Y):
    def _update_cat_count_dims(cat_count, highest_feature):
        diff = highest_feature + 1 - cat_count.shape[1]
        if diff > 0:
            # we append a column full of zeros for each new category
            return np.pad(cat_count, [(0, 0), (0, diff)], "constant")
        return cat_count

    def _update_cat_count(X_feature, Y, cat_count, n_classes):
        for j in range(n_classes):
            mask = Y[:, j].astype(bool)
            if Y.dtype.type == np.int64:
                weights = None
            else:
                weights = Y[mask, j]
            counts = np.bincount(X_feature[mask], weights=weights)
            indices = np.nonzero(counts)[0]
            cat_count[j, indices] += counts[indices]

    self.class_count_ += Y.sum(axis=0)
    self.n_categories_ = self._validate_n_categories(X, self.min_categories)
    for i in range(self.n_features_):
        X_feature = X[:, i]
        self.category_count_[i] = _update_cat_count_dims(
            self.category_count_[i], self.n_categories_[i] - 1
        )
        _update_cat_count(
            X_feature, Y, self.category_count_[i], self.class_count_.shape[0]
        )