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 module

selective_scan_ops

Selective scan operation registrations for Mamba SSM.

This module registers the following ops:

• selective_scan_fwd: Full selective scan forward pass
• selective_scan_fwd_minimal: Minimal variant without optional tensors
• selective_scan_update: Single-step update for autoregressive inference

Structs

• ​SelectiveScanFwd: Selective scan forward pass operation for Mamba SSM.
• ​SelectiveScanFwdMinimal: Minimal selective scan forward pass - no optional D, z, or delta_bias.
• ​SelectiveScanUpdate: Selective scan update operation for autoregressive inference.

Functions

• ​selective_scan_fwd_minimal_shape: Returns the output shape for the selective_scan_fwd_minimal op.
• ​selective_scan_fwd_shape: Returns the output shape for the selective_scan_fwd op.
• ​selective_scan_update_shape: Returns the output shapes for the selective_scan_update op.