Debug a NumPy Extremely Randomized Trees Classifier

Context

Candidates for the Anthropic Fellows program report a second online assessment that is a debugging task: a NumPy implementation of Extremely Randomized Trees (Extra-Trees) that crashes on some datasets and scores far below a reference library on others. No AI tools are allowed. Below is a self-contained practice version written for this page; it is not the original assessment code.

Task

The module below implements an Extra-Trees classifier from scratch. The test suite reports:

• ValueError: setting an array element with a sequence ... inhomogeneous shape on the breast-cancer dataset and on any dataset with a constant column.
• On a 3-class synthetic dataset, predict returns an array of the wrong shape.

Fix every bug you can find so that:

1. It never crashes on numeric data with any label type (ints, strings), including constant features.
2. predict(X) returns one label per row.
3. Test accuracy is within a few points of sklearn.ensemble.ExtraTreesClassifier(n_estimators=50), which gets about 92% on the synthetic set and 94% on breast cancer.

Do not rewrite the algorithm or change the public API. Extra-Trees means: at each node, pick max_features random features, draw one uniform random threshold per feature between that feature's min and max in the node, keep the split with the lowest weighted Gini impurity, and recurse. Each tree is trained on the full training set; the forest averages per-tree class probabilities.

Buggy Code

import numpy as np


class Node:
    def __init__(self, feature=None, threshold=None, left=None, right=None, proba=None):
        self.feature, self.threshold = feature, threshold
        self.left, self.right, self.proba = left, right, proba


def gini(y, n_classes):
    counts = np.bincount(y, minlength=n_classes)
    return 1.0 - np.sum(counts ** 2)


class ExtraTree:
    def __init__(self, max_features, min_samples_split=2, max_depth=None, seed=0):
        self.max_features = max_features
        self.min_samples_split = min_samples_split
        self.max_depth = max_depth
        self.rng = np.random.default_rng(seed)

    def fit(self, X, y, n_classes):
        self.n_classes = n_classes
        self.root = self._build(X, y, depth=0)
        return self

    def _leaf(self, y):
        counts = np.bincount(y)
        return Node(proba=counts / counts.sum())

    def _build(self, X, y, depth):
        if (len(y) < self.min_samples_split
                or (self.max_depth is not None and depth >= self.max_depth)
                or np.all(y == y[0])):
            return self._leaf(y)
        features = self.rng.choice(X.shape[1], self.max_features)
        best = None
        for f in features:
            lo, hi = X[:, f].min(), X[:, f].max()
            t = self.rng.uniform(lo, hi)
            left = X[:, f] < t
            n_l = left.sum()
            n_r = len(y) - n_l
            score = (n_l * gini(y[left], self.n_classes)
                     + n_r * gini(y[~left], self.n_classes)) / len(y)
            if best is None or score > best[0]:
                best = (score, f, t, left)
        _, f, t, left = best
        return Node(f, t,
                    self._build(X[left], y[left], depth + 1),
                    self._build(X[~left], y[~left], depth + 1))

    def predict_proba_one(self, x):
        node = self.root
        while node.proba is None:
            node = node.right if x[node.feature] < node.threshold else node.left
        return node.proba


class ExtraTreesClassifier:
    def __init__(self, n_estimators=50, max_features="sqrt", seed=0, **tree_kw):
        self.n_estimators = n_estimators
        self.max_features = max_features
        self.seed = seed
        self.tree_kw = tree_kw

    def fit(self, X, y):
        X = np.asarray(X, dtype=float)
        y = np.asarray(y)
        self.classes_, y_enc = np.unique(y, return_inverse=True)
        d = X.shape[1]
        k = max(1, int(np.sqrt(d))) if self.max_features == "sqrt" else self.max_features
        self.trees = [
            ExtraTree(k, seed=self.seed, **self.tree_kw).fit(X, y_enc, len(self.classes_))
            for _ in range(self.n_estimators)
        ]
        return self

    def predict_proba(self, X):
        X = np.asarray(X, dtype=float)
        per_tree = [[t.predict_proba_one(x) for x in X] for t in self.trees]
        return np.mean(per_tree, axis=1)

    def predict(self, X):
        return self.classes_[np.argmax(self.predict_proba(X), axis=1)]

What to Submit

The corrected module, plus a short list of each bug: the symptom it caused and your fix.