from sklearn.datasets import load_digits
from sklearn.kernel_approximation import AdditiveChi2Sampler

X, y = load_digits(return_X_y=True)
n_features = X.shape[1]

for sample_steps in range(1, 4):
    
    chi2sampler = AdditiveChi2Sampler(sample_steps=sample_steps)

    # Updated formula
    n_features_out = n_features * (2 * sample_steps - 1)
    X_transformed = chi2sampler.fit_transform(X, y)

    assert X_transformed.shape[1] == n_features_out