For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).
Mojo struct
Struct_gumbel_argmax_from_probs
struct Struct_gumbel_argmax_from_probs
Registers the sampler.gumbel_argmax_from_probs graph op.
Draws one token per row proportionally to a row of unnormalized
probabilities, by Gumbel-max over ln(p). The noise comes from the
per-row seed inside the kernel, so the caller passes no noise tensor.
Rows with equal seeds draw with equal noise, which is how a request's
draft positions share one noise row.
Implemented traitsโ
Methodsโ
executeโ
static def execute[target: StringSpan[ImmStaticOrigin], _trace_name: StringSpan[ImmStaticOrigin]](out_tokens: ManagedTensorSlice[IOSpec[_, _].Output, static_spec=out_tokens.static_spec], seed: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=seed.static_spec], probs: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=probs.static_spec], ctx: DeviceContext)