def create_vit_basic_head(
    *,
    # Projection configs.
    in_features: int,
    out_features: int,
    # Pooling configs.
    seq_pool_type: str = "cls",
    # Dropout configs.
    dropout_rate: float = 0.5,
    # Activation configs.
    activation: Callable = None,
) -> Module:
    """
    Creates vision transformer basic head.

    ::


                                        Pooling
                                           ↓
                                        Dropout
                                           ↓
                                       Projection
                                           ↓
                                       Activation


    Activation examples include: ReLU, Softmax, Sigmoid, and None.
    Pool type examples include: cls, mean and none.

    Args:

        in_features: input channel size of the resnet head.
        out_features: output channel size of the resnet head.

        pool_type (str): Pooling type. It supports "cls", "mean " and "none". If set to
            "cls", it assumes the first element in the input is the cls token and
            returns it. If set to "mean", it returns the mean of the entire sequence.

        activation (callable): a callable that constructs vision transformer head
            activation layer, examples include: nn.ReLU, nn.Softmax, nn.Sigmoid, and
            None (not applying activation).

        dropout_rate (float): dropout rate.
    """
    assert seq_pool_type in ["cls", "mean", "none"]

    if seq_pool_type in ["cls", "mean"]:
        seq_pool_model = SequencePool(seq_pool_type)
    elif seq_pool_type == "none":
        seq_pool_model = None
    else:
        raise NotImplementedError

    if activation is None:
        activation_model = None
    elif activation == Softmax:
        activation_model = activation(dim=1)
    else:
        activation_model = activation()

    return VisionTransformerBasicHead(
        sequence_pool=seq_pool_model,
        dropout=Dropout(dropout_rate) if dropout_rate > 0.0 else None,
        proj=Linear(in_features, out_features),
        activation=activation_model,
    )
Creates vision transformer basic head.

::


Pooling
↓
Dropout
↓
Projection
↓
Activation


Activation examples include: ReLU, Softmax, Sigmoid, and None.
Pool type examples include: cls, mean and none.

Args:

in_features: input channel size of the resnet head.
out_features: output channel size of the resnet head.

pool_type (str): Pooling type. It supports "cls", "mean " and "none". If set to
"cls", it assumes the first element in the input is the cls token and
returns it. If set to "mean", it returns the mean of the entire sequence.

activation (callable): a callable that constructs vision transformer head
activation layer, examples include: nn.ReLU, nn.Softmax, nn.Sigmoid, and
None (not applying activation).

dropout_rate (float): dropout rate.
python/aitemplate/frontend/nn/head.py