[builder.py] How to Expand an input to a shape which is determined by another input? i.e., 'shape' is determined at runtime

[kbiscoding]

Hello,
I am trying to add Qwen2.5VL-3B > Text Rotary embedding logic to my copy of onnxruntime-genai > python/py/models/builder.py

I am trying to use node "Expand" for initializing rotary sin-cos embeddings for Qwen2.5VL-3B.
The 'shape' input to Expand is made from position_ids.shape[1], where position_ids is an input and does not have a fixed shape.

See python3.11/site-packages/transformers/models/qwen2_5_vl/modeling_qwen2_5_vl.py > Qwen2_5_VLRotaryEmbedding > forward:
inv_freq_expanded = self.inv_freq[None, None, :, None].float().expand(3, position_ids.shape[1], -1, 1)

So, in my copy of builder.py I tried to make a "Shape" node to obtain shape[1]
self.make_shape("/model/rot_emb_cache_from_scratch/shape_of_position_ids", position_ids, shape=[1])

This produces output node '/model/rot_emb_cache_from_scratch/shape_of_position_ids/output' providing value for position_ids.shape[1]

Now, how do I make this value as an input to "Expand" node?

As mentioned above, the transformers code I am trying to mirror is:
inv_freq_expanded = inv_freq[None, None, :, None].float().expand(3, position_ids.shape[1], -1, 1)

So I do
reshaped_inv_freq_for_expansion = inv_freq[None, None, :, None].float()
self.make_initializer(reshaped_inv_freq_for_expansion,
'/model/rot_emb_cache_from_scratch/inv_freq_before_expansion',
self.io_dtype)
self.make_expand('/model/rot_emb_cache_from_scratch/inv_freq_expand',
['/model/rot_emb_cache_from_scratch/inv_freq_before_expansion'],
self.io_dtype, shape=[3, ' HOW DO I PASS VALUE HERE?', -1, 1])

How do I make a tensor from the output '/model/rot_emb_cache_from_scratch/shape_of_position_ids/output' and pass it to Expand node as shape?

I tried searching online and asking code assistants, it seems like we cannot attach a tensor to ONNX graph if the value is determined at runtime. Any workarounds?

[kunal-vaishnavi]

Thank you for your efforts!

The line, position_ids.shape[1], can be translated to position_ids --> Shape --> Gather (idx = 1). Following the output naming structure, the output of the Gather node will be f"{gather_node_name}/output_0". That output is the name of the value to pass to the Expand node. The shape input for the Expand node would then be [3, f"{gather_op_name}/output_0", -1, 1]. The shape of the Expand node would be the output shape that results from the operation (e.g. ['batch_size', 'sequence_length', hidden_size]).

One note is that some of the rotary embedding logic can usually be pre-computed ahead of time. For example, the make_rotary_embedding_caches_from_scratch function pre-computes the rotary embedding caches up to the max sequence length that the model can take. I have not looked at Qwen-2.5-VL but you can investigate if you can override this method and bring over the PyTorch logic in your Qwen-2.5-VL class implementation instead. You can look at how OpenAI's gpt-oss model is implemented as an example.