bayesianbandits.ArmFeaturizer

class bayesianbandits.ArmFeaturizer

Bases: ABC, Generic[TokenType]

Abstract base class for vectorized arm feature transformation.

Transforms features for multiple arms in a single call to enable efficient batched processing in shared model bandits.

Type Parameters

TokenTypetype

The type of action tokens this featurizer accepts.

abstract transform(X: Sized, *, action_tokens: Sequence[TokenType]) → Any

Transform features for all arms in a single vectorized call.

Parameters:

• X (Sized) – Input context features. Implementations typically expect array-like of shape (n_contexts, n_features).

• action_tokens (sequence of TokenType, length n_arms) – Action tokens for each arm to featurize.

Returns:

X_transformed – Stacked features where rows are ordered as: - Rows 0:n_contexts are for action_tokens[0] - Rows n_contexts:2*n_contexts are for action_tokens[1] - etc.

The specific type depends on the implementation and input type.

Return type:

array-like of shape (n_contexts * n_arms, n_features_out)

Notes

The output is “stacked” such that each context is repeated for each arm, with arm-specific features appended. This allows a single model to process all arm-context pairs efficiently.

Examples

>>> import numpy as np
>>> # Discrete action tokens
>>> X = np.array([[1, 2], [3, 4]])  # 2 contexts, 2 features
>>> action_tokens = [0, 1, 2]  # 3 arms
>>> # After transform: 6 rows (2 contexts * 3 arms)
>>> # Rows 0-1: context features + arm 0 features
>>> # Rows 2-3: context features + arm 1 features
>>> # Rows 4-5: context features + arm 2 features