IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /get-started.md). For the complete documentation index, see llms.txt.
Skip to main content
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โ€‹

AnyType, Deinitable, Movable

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)