GatherOperationsConverter

Converter for memory-access / indexing operations.

Today that's just gather and its framework aliases — the critical path for LLM exports, where every transformer forward pass begins with a token-id \u2192 embedding lookup. Without a converter for gather / embedding / index_select, a traced Llama / Mistral / Qwen / Gemma model fails at the very first operation and never reaches the norms, activations, or attention that the other P1 converters cover.

The target lowering is the canonical embedding(input_ids) shape: a 1-D index tensor indexing the leading dimension of a 2-D embedding weight. Higher-rank gathers (attention-side index gathers, multi-dim scatter/gather) can be added in follow-up PRs once a traced model surfaces them; scoping this converter to the LLM front-door case keeps review tight.

Emitted shape:

%out = stablehlo.gather(%weights, %indices)
{ dimension_numbers = #stablehlo.gather<
offset_dims = [1],
collapsed_slice_dims = [0],
start_index_map = [0],
index_vector_dim = 1>,
slice_sizes = array<i64: 1, hidden_size>,
indices_are_sorted = false }
: (tensor<vocab_size x hidden_size x f32>, tensor<seq_len x i32>)
-> tensor<seq_len x hidden_size x f32>

The slice_sizes vector is derived from the weight shape: a 1 along the gathered axis and the full extent of every other dimension. offset_dims, collapsed_slice_dims, and start_index_map are computed from the single gather axis.

Constructors

Link copied to clipboard
constructor()

Properties

Link copied to clipboard
open override val supportedOperations: Set<String>

Set of operation names this converter supports

Functions

Link copied to clipboard
open override fun convert(node: GraphNode, operands: List<String>, context: ConversionContext): ConversionResult

Convert a graph node to StableHLO operations