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.