flashinfer.logits_processor.TaggedTensor

class flashinfer.logits_processor.TaggedTensor(data: Tensor, type: TensorType)

一个张量包装器，在流水线执行过程中维护语义类型信息。

它确保了类型安全，因为张量流经 logits 处理流水线，并且对下游 PyTorch 代码没有任何摩擦

注意

• TaggedTensor 主要供 LogitsPipe 内部使用

• 用户通常使用普通张量；标记是自动发生的

__init__(data: Tensor, type: TensorType) → None

方法

__init__(data, type)
indices(t)	创建一个类型为 TensorType.INDICES 的 TaggedTensor。
logits(t)	创建一个类型为 TensorType.LOGITS 的 TaggedTensor。
probs(t)	创建一个类型为 TensorType.PROBS 的 TaggedTensor。
size([dim])	获取底层张量的尺寸。

属性

data	底层的张量。
type	张量的语义类型。
device	获取底层张量的设备。
dtype	获取底层张量的数据类型。
shape	获取底层张量的形状。