flashinfer.logits_processor.Softmax

class flashinfer.logits_processor.Softmax(**params: Any)

Softmax processor to convert logits to probabilities.

Applies the softmax function.

TensorType.LOGITS -> TensorType.PROBS

Examples

>>> import torch
>>> from flashinfer.logits_processor import LogitsPipe, Softmax, Sample
>>> torch.manual_seed(42)
>>> pipe = LogitsPipe([Softmax()])
>>> logits = torch.randn(2, 2, device="cuda")
>>> logits
tensor([[ 0.1940,  2.1614], [ -0.1721,  0.8491]], device='cuda:0')
>>> probs = pipe(logits)
>>> probs
tensor([[0.1227, 0.8773], [0.2648, 0.7352]], device='cuda:0')

Notes

Can only appear once in a pipeline.

__init__(**params: Any)

Initialize the processor.

Parameters:

**params (Any) – Processor-specific parameters at compile-time.

Methods

__init__(**params)	Initialize the processor.
legalize(input_type)	Legalize the processor into a list of low-level operators.